Update README and fix a few TF2 compatibility issues.

PiperOrigin-RevId: 289177692

scripts/tf_cnn_benchmarks/README.md

@@ -1,18 +1,16 @@
 # tf_cnn_benchmarks: High performance benchmarks
 
-tf_cnn_benchmarks contains implementations of several popular convolutional
-models, and is designed to be as fast as possible. tf_cnn_benchmarks supports
-both running on a single machine or running in distributed mode across multiple
-hosts. See the [High-Performance models
-guide](https://www.tensorflow.org/performance/performance_models) for more
-information.
-
-These models utilize many of the strategies in the [TensorFlow Performance
-Guide](https://www.tensorflow.org/performance/performance_guide). Benchmark
-results can be found [here](https://www.tensorflow.org/performance/benchmarks).
-
-These models are designed for performance. For models that have clean and
-easy-to-read implementations, see the [TensorFlow Official
+**Note: tf_cnn_benchmarks is no longer maintained.**
+
+tf_cnn_benchmarks contains TensorFlow 1 implementations of several popular
+convolutional models, and is designed to be as fast as possible.
+tf_cnn_benchmarks supports both running on a single machine or running in
+distributed mode across multiple hosts.
+
+tf_cnn_benchmarks is no longer maintained. Although it will run with TensorFlow
+2, it was written and optimized for TensorFlow 1, and has not been maintained
+since TensorFlow 2 was released. For clean and easy-to-read TensorFlow 2 models,
+please see the [TensorFlow Official
 Models](https://github.com/tensorflow/models/tree/master/official).
 
 ## Getting Started
@@ -23,14 +21,15 @@ To run ResNet50 with synthetic data without distortions with a single GPU, run
 python tf_cnn_benchmarks.py --num_gpus=1 --batch_size=32 --model=resnet50 --variable_update=parameter_server
 ```
 
-Note that the master branch of tf_cnn_benchmarks requires the latest nightly
-version of TensorFlow. You can install the nightly version by running `pip
-install tf-nightly-gpu` in a clean environment, or by installing TensorFlow from
-source. We sometimes will create a branch of tf_cnn_benchmarks, in the form of
-cnn_tf_vX.Y_compatible, that is compatible with TensorFlow version X.Y For
-example, branch
+Note that the master branch of tf_cnn_benchmarks occasionally requires the
+latest nightly version of TensorFlow. You can install the nightly version by
+running `pip install tf-nightly-gpu` in a clean environment, or by installing
+TensorFlow from source. We sometimes will create a branch of tf_cnn_benchmarks,
+in the form of cnn_tf_vX.Y_compatible, that is compatible with TensorFlow
+version X.Y. For example, branch
 [cnn_tf_v1.9_compatible](https://github.com/tensorflow/benchmarks/tree/cnn_tf_v1.9_compatible/scripts/tf_cnn_benchmarks)
-works with TensorFlow 1.9.
+works with TensorFlow 1.9. However, as tf_cnn_benchmarks is no longer
+maintained, we will likely no longer create new branches.
 
 Some important flags are
 

scripts/tf_cnn_benchmarks/benchmark_cnn_test.py

@@ -704,7 +704,7 @@ def testMoveTrainDir(self):
     params = params._replace(train_dir=new_train_dir, eval=True)
     self._run_benchmark_cnn_with_black_and_white_images(params)
 
-  @mock.patch('tensorflow.train.Saver')
+  @mock.patch('tensorflow.compat.v1.train.Saver')
   @mock.patch('benchmark_cnn._get_checkpoint_to_load')
   def testLoadCheckpoint(self, mock_checkpoint_to_load, mock_saver):
     """Tests load checkpoint with full path to checkpoint."""

scripts/tf_cnn_benchmarks/run_tests.py

@@ -23,7 +23,7 @@
 
 from absl import app
 from absl import flags as absl_flags
-import tensorflow.compat.v2 as tf
+import tensorflow.compat.v1 as tf
 
 import all_reduce_benchmark_test
 import allreduce_test
@@ -32,7 +32,6 @@
 import cnn_util_test
 import variable_mgr_util_test
 from models import model_config
-from models.tf1_only import nasnet_test
 
 # Ideally, we wouldn't need this option, and run both distributed tests and non-
 # distributed tests. But, TensorFlow allocates all the GPU memory by default, so
@@ -67,6 +66,7 @@ def main(_):
         loader.loadTestsFromModule(all_reduce_benchmark_test),
     ])
     if model_config.can_import_contrib:
+      from models.tf1_only import nasnet_test  # pylint: disable=g-import-not-at-top
       suite.addTest(loader.loadTestsFromModule(nasnet_test))
     dist_suite = unittest.TestSuite([
         loader.loadTestsFromModule(benchmark_cnn_distributed_test),