Squash commit.

PiperOrigin-RevId: 517247440

scripts/tf_cnn_benchmarks/all_reduce_benchmark.py

@@ -24,11 +24,6 @@
 TODO(reedwm): Allow custom sizes to be specified.
 """
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-
 import os
 import time
 
@@ -43,14 +38,14 @@
 from cnn_util import log_fn
 
 
-absl_flags.DEFINE_integer('iters_per_step', 5,
-                          'Number of iterations to run all-reduce for, per '
-                          'step. Every step, a session will be run on a Graph '
-                          'that contains this many copies of the all-reduce. '
-                          'The copies are run sequentially. Setting this above '
-                          '1 is useful to lower the overhead of starting the '
-                          'session run, running the VariableV2 ops at the '
-                          'start of the step, etc.')
+_ITERS_PER_STEP = absl_flags.DEFINE_integer(
+    'iters_per_step', 5, 'Number of iterations to run all-reduce for, per '
+    'step. Every step, a session will be run on a Graph '
+    'that contains this many copies of the all-reduce. '
+    'The copies are run sequentially. Setting this above '
+    '1 is useful to lower the overhead of starting the '
+    'session run, running the VariableV2 ops at the '
+    'start of the step, etc.')
 
 
 flags.define_flags()
@@ -283,7 +278,7 @@ def main(positional_arguments):
   tfversion = cnn_util.tensorflow_version_tuple()
   log_fn('TensorFlow:  %i.%i' % (tfversion[0], tfversion[1]))
 
-  run_benchmark(bench, absl_flags.FLAGS.iters_per_step)
+  run_benchmark(bench, _ITERS_PER_STEP.value)
 
 if __name__ == '__main__':
   tf.disable_v2_behavior()

scripts/tf_cnn_benchmarks/all_reduce_benchmark_test.py

@@ -14,10 +14,6 @@
 # ==============================================================================
 """Tests for all_reduce_benchmark.py."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import tensorflow.compat.v1 as tf
 
 import all_reduce_benchmark

scripts/tf_cnn_benchmarks/allreduce.py

@@ -14,14 +14,9 @@
 # ==============================================================================
 """Utilities for allreduce."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import collections as pycoll
 import re
 
-from six.moves import xrange  # pylint: disable=redefined-builtin
 import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
@@ -451,7 +446,7 @@ def sum_gradients_all_reduce(single_session,
   gv = list(zip(*tower_grads))
   merge_scope = allreduce_merge_scope if allreduce_merge_scope > 0 else 1
   chunked_gv = [gv[x:x + merge_scope]
-                for x in xrange(0, len(gv), merge_scope)]
+                for x in range(0, len(gv), merge_scope)]
   for chunk in chunked_gv:
     with tf.name_scope('allreduce'):
       for grad_and_vars in chunk:
@@ -638,7 +633,7 @@ def unpack_small_tensors(tower_grads, packing):
   num_packed = len(packing.keys()) // num_devices
   for dev_idx, gv_list in enumerate(tower_grads):
     new_gv_list = gv_list[num_packed:]
-    for i in xrange(0, num_packed):
+    for i in range(0, num_packed):
       k = '%d:%d' % (dev_idx, i)
       gpt = packing[k]
       gv = unpack_grad_tuple(gv_list[i], gpt)

scripts/tf_cnn_benchmarks/allreduce_test.py

@@ -15,10 +15,6 @@
 
 """Tests for tf_cnn_benchmark.allreduce."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import collections as pycoll
 
 import numpy as np

scripts/tf_cnn_benchmarks/batch_allreduce.py

@@ -24,18 +24,13 @@
 across devices and inter-device tensor transfers across device links.
 """
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 # TODO(reedwm): Support distributed all-reduces in this file.
 # TODO(reedwm): Merge this code with allreduce.py, which contains some batch
 # all-reduce code that this file calls. allreduce.py also supports distributed
 # batch-reduce while this file only supports single-machine all-reduce.
 
 import abc
 
-import six
 import tensorflow.compat.v1 as tf
 
 from tensorflow.python.ops import data_flow_ops
@@ -60,8 +55,7 @@ def _all_reduce_using_copy(tensors_across_devices, use_mean):
   return reduced_tensor
 
 
-@six.add_metaclass(abc.ABCMeta)
-class BatchAllReduceAlgorithm(object):
+class BatchAllReduceAlgorithm(metaclass=abc.ABCMeta):
   """Represents an algorithm for performing a batch all-reduce operation."""
 
   def batch_all_reduce(self,

scripts/tf_cnn_benchmarks/benchmark_cnn.py

@@ -17,10 +17,6 @@
 See the README for more information.
 """
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import argparse
 from collections import namedtuple
 import contextlib
@@ -34,9 +30,6 @@
 
 from absl import flags as absl_flags
 import numpy as np
-
-import six
-from six.moves import xrange  # pylint: disable=redefined-builtin
 import tensorflow.compat.v1 as tf
 
 # pylint: disable=g-direct-tensorflow-import
@@ -54,7 +47,7 @@
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python import debug as tf_debug
 from tensorflow.python.client import timeline
-from tensorflow.python.framework import graph_util
+from tensorflow.python.framework import convert_to_constants
 from tensorflow.python.framework import graph_util_impl
 from tensorflow.python.framework import importer
 from tensorflow.python.ops import data_flow_ops
@@ -1526,7 +1519,7 @@ def __init__(self, params, dataset=None, model=None):
     # compute device, and never on a parameter server device.
     self.raw_devices = [
         '%s/%s:%i' % (worker_prefix, self.params.device, i)
-        for i in xrange(self.num_gpus)
+        for i in range(self.num_gpus)
     ]
 
     subset = 'validation' if params.eval else 'train'
@@ -1747,7 +1740,7 @@ def reset_devices_for_task(self, task_num, is_local=False):
     self.cpu_device = '%s/cpu:0' % worker_prefix
     self.raw_devices = [
         '%s/%s:%i' % (worker_prefix, self.params.device, i)
-        for i in xrange(self.num_gpus)
+        for i in range(self.num_gpus)
     ]
     self.devices = self.variable_mgr.get_devices()
 
@@ -1759,8 +1752,8 @@ def raw_devices_across_tasks(self, is_local=False):
     else:
       return [
           'job:worker/replica:0/task%s/%s:%i' % (t, self.params.device, i)
-          for t in xrange(self.num_workers)
-          for i in xrange(self.num_gpus)
+          for t in range(self.num_workers)
+          for i in range(self.num_gpus)
       ]
 
   def print_info(self):
@@ -1985,7 +1978,7 @@ def _initialize_eval_graph(self, enqueue_ops, input_producer_op,
             self.params.use_python32_barrier)
         image_producer.start()
       if enqueue_ops:
-        for i in xrange(len(enqueue_ops)):
+        for i in range(len(enqueue_ops)):
           sess.run(enqueue_ops[:(i + 1)])
           if image_producer is not None:
             image_producer.notify_image_consumption()
@@ -2003,7 +1996,7 @@ def _eval_once(self, sess, summary_writer, fetches, summary_op,
       top_1_accuracy_sum = 0.0
       top_5_accuracy_sum = 0.0
       total_eval_count = self.num_batches * self.batch_size
-      for step in xrange(self.num_batches):
+      for step in range(self.num_batches):
         if (summary_writer and self.params.save_summaries_steps > 0 and
             (step + 1) % self.params.save_summaries_steps == 0):
           results, summary_str = sess.run([fetches, summary_op])
@@ -2339,7 +2332,7 @@ def benchmark_with_session(self, sess, supervisor, graph_info,
           self.params.use_python32_barrier)
       image_producer.start()
     if graph_info.enqueue_ops:
-      for i in xrange(len(graph_info.enqueue_ops)):
+      for i in range(len(graph_info.enqueue_ops)):
         sess.run(graph_info.enqueue_ops[:(i + 1)])
         if image_producer is not None:
           image_producer.notify_image_consumption()
@@ -2617,7 +2610,7 @@ def _preprocess_graph(self, graph, graph_info):
       with tf.Session(config=create_config_proto(self.params)) as sess:
         sess.run(tf.global_variables_initializer())
         sess.run(tf.local_variables_initializer())
-        graphdef = graph_util.convert_variables_to_constants(
+        graphdef = convert_to_constants.convert_variables_to_constants(
             sess,
             graphdef,
             output_node_names,
@@ -2749,7 +2742,7 @@ def _build_input_processing(self, shift_ratio=0):
             shared_name='input_producer_staging_area_%d_eval_%s' %
             (device_num, self._doing_eval))
         input_producer_stages.append(staging_area)
-        for group_index in xrange(self.batch_group_size):
+        for group_index in range(self.batch_group_size):
           batch_index = group_index + device_num * self.batch_group_size
           put_op = staging_area.put(
               [parts[batch_index] for parts in input_list])
@@ -2868,7 +2861,8 @@ def _build_model(self):
       if self.variable_mgr.supports_staged_vars():
         for staging_ops in self.variable_mgr.staging_vars_on_devices:
           gpu_compute_stage_ops.extend(
-              [put_op for _, (put_op, _) in six.iteritems(staging_ops)])
+              [put_op for _, (put_op, _) in staging_ops.items()]
+          )
       enqueue_ops.append(tf.group(*gpu_compute_stage_ops,
                                   name='gpu_compute_stage_ops_group'))
       if gpu_grad_stage_ops:

scripts/tf_cnn_benchmarks/benchmark_cnn_distributed_test.py

@@ -20,10 +20,6 @@
 The output for each process is written to disk and can be viewed to debug tests.
 See get_test_output_dir() in platforms/default/util.py for more info.
 """
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
 from collections import namedtuple
 import os
 import subprocess
@@ -32,7 +28,6 @@
 
 from absl import flags as absl_flags
 import portpicker
-import six
 import tensorflow.compat.v1 as tf
 import flags
 import test_util
@@ -55,7 +50,8 @@ def _convert_params_to_flags_list(params):
     A list of flags.
   """
   return [
-      '--%s=%s' % (k, str(v)) for k, v in six.iteritems(params._asdict())
+      '--%s=%s' % (k, str(v))
+      for k, v in params._asdict().items()
       if v != flags.param_specs[k].default_value
   ]
 

scripts/tf_cnn_benchmarks/benchmark_cnn_distributed_test_runner.py

@@ -21,10 +21,6 @@
 file is spawned as a Python process, which supports the added functionality.
 """
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from absl import flags as absl_flags
 import numpy as np
 import tensorflow.compat.v1 as tf
@@ -34,8 +30,8 @@
 import test_util
 
 
-absl_flags.DEFINE_string('fake_input', 'none',
-                         """What fake input to inject into benchmark_cnn. This
+_FAKE_INPUT = absl_flags.DEFINE_string(
+    'fake_input', 'none', """What fake input to inject into benchmark_cnn. This
                             is ignored if --model=test_model.
                             Options are:
                             none: Do not inject any fake input.
@@ -44,7 +40,6 @@
                             label of 1.""")
 
 flags.define_flags()
-FLAGS = flags.FLAGS
 
 
 def get_test_image_preprocessor(batch_size, params):
@@ -60,9 +55,9 @@ def get_test_image_preprocessor(batch_size, params):
   Raises:
     ValueError: Flag --fake_input is an invalid value.
   """
-  if FLAGS.fake_input == 'none':
+  if _FAKE_INPUT.value == 'none':
     return None
-  elif FLAGS.fake_input == 'zeros_and_ones':
+  elif _FAKE_INPUT.value == 'zeros_and_ones':
     half_batch_size = batch_size // 2
     images = np.zeros((batch_size, 227, 227, 3), dtype=np.float32)
     images[half_batch_size:, :, :, :] = 1
@@ -75,7 +70,7 @@ def get_test_image_preprocessor(batch_size, params):
     preprocessor.expected_subset = 'validation' if params.eval else 'train'
     return preprocessor
   else:
-    raise ValueError('Invalid --fake_input: %s' % FLAGS.fake_input)
+    raise ValueError('Invalid --fake_input: %s' % _FAKE_INPUT.value)
 
 
 def run_with_real_model(params):

scripts/tf_cnn_benchmarks/benchmark_cnn_test.py

@@ -14,10 +14,6 @@
 # ==============================================================================
 
 """Tests for benchmark_cnn."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
 import glob
 import os
 import re

scripts/tf_cnn_benchmarks/cnn_util.py

@@ -14,9 +14,6 @@
 # ==============================================================================
 
 """Utilities for CNN benchmarks."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
 
 import sys
 import threading

scripts/tf_cnn_benchmarks/cnn_util_test.py

@@ -15,10 +15,6 @@
 
 """Tests for tf_cnn_benchmarks.cnn_util."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import threading
 import time
 

scripts/tf_cnn_benchmarks/coco_metric.py

@@ -19,20 +19,14 @@
 COCO API: github.com/cocodataset/cocoapi/
 """
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import atexit
 import tempfile
 
 from absl import flags
-
 import numpy as np
+import pycocotools.coco
 from pycocotools.coco import COCO
 from pycocotools.cocoeval import COCOeval
-import six
-
 import tensorflow.compat.v1 as tf
 
 import mlperf
@@ -42,9 +36,8 @@
 
 
 # https://github.com/cocodataset/cocoapi/issues/49
-if six.PY3:
-  import pycocotools.coco
-  pycocotools.coco.unicode = str
+
+pycocotools.coco.unicode = str
 
 
 def async_eval_runner(queue_predictions, queue_results, val_json_file):

scripts/tf_cnn_benchmarks/constants.py

@@ -14,10 +14,6 @@
 # ==============================================================================
 """Constants used in tf_cnn_benchmarks."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 from enum import Enum
 
 # Results fetched with this prefix will not be reduced. Instead, they will be