<NME> benchmark_cnn.py <BEF> # Copyright 2017 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """TensorFlow benchmark library. 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 import math import multiprocessing import os import re import threading import time import traceback 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 import cnn_util import constants import datasets import flags import mlperf import variable_mgr import variable_mgr_util from cnn_util import log_fn from models import model_config from platforms import util as platforms_util from google.protobuf import text_format 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 graph_util_impl from tensorflow.python.framework import importer from tensorflow.python.ops import data_flow_ops from tensorflow.python.platform import gfile from tensorflow.python.util import nest _DEFAULT_NUM_BATCHES = 100 # GraphInfo encapsulates the tensors/ops that we care about after building a # graph. We use them to benchmark the graph. GraphInfo = namedtuple( # pylint: disable=invalid-name 'GraphInfo', [ # Ops that produce the input batches (before preprocessing). 'input_producer_op', # Ops that adds the preprocessed images to the staging areas 'enqueue_ops', # Fetches of sess.run() 'fetches', # Op that performs synchronization in distributed mode 'execution_barrier', # The global step variable 'global_step', # Group of ops that perform per-device initialization work 'local_var_init_op_group', # Op to produce summaries 'summary_op' ]) # InputProcessingInfo contains various sources of inputs which will be later fed # into the model. If synthetic data is used, all three fields are None. InputProcessingInfo = namedtuple( 'InputProcessingInfo', [ # The first two fields are non-None iff datasets prefetching is not # used. # Ops that produce the input batches. 'input_producer_op', # A list of StagingArea for each device. 'input_producer_stages', # Input produced using multi device iterator. Non-None iff datasets # prefetching is used 'multi_device_iterator_input' ]) # TODO(reedwm): add upper_bound and lower_bound to appropriate integer and # float flags, and change certain string flags to enum flags. flags.DEFINE_string('model', 'trivial', 'Name of the model to run, the list of supported models ' 'are defined in models/model.py') # The code will first check if it's running under benchmarking mode # or evaluation mode, depending on 'eval': # Under the evaluation mode, this script will read a saved model, # and compute the accuracy of the model against a validation dataset. # Additional ops for accuracy and top_k predictors are only used under # this mode. # Under the benchmarking mode, user can specify whether nor not to use # the forward-only option, which will only compute the loss function. # forward-only cannot be enabled with eval at the same time. flags.DEFINE_boolean('eval', False, 'whether use eval or benchmarking') flags.DEFINE_integer('eval_interval_secs', 0, 'How often to run eval on saved checkpoints. Usually the ' 'same as save_model_secs from the corresponding training ' 'run. Pass 0 to eval only once.') flags.DEFINE_integer('eval_during_training_every_n_steps', None, 'Every n steps during training, pause training, run ' 'evaluation, then resume training. Must not be used with ' '--eval, as unlike --eval, this option causes both ' 'training and eval to be done. This may take slightly ' 'more GPU memory than running just training or evaluation ' 'alone. It also may slightly slow down training, even ' 'when not taking into account the additional time to ' 'evaluate.', lower_bound=1) flags.DEFINE_float('eval_during_training_every_n_epochs', None, 'After every n training epochs, pause training, run ' 'evaluation, then resume training. See ' '--eval_during_training_every_n_steps for more information.') flags.DEFINE_list('eval_during_training_at_specified_steps', [], 'Specify a list of training steps, pause training at each of ' 'these steps, run evaluation, then resume training. See ' '--eval_during_training_every_n_steps for more information.') flags.DEFINE_list('eval_during_training_at_specified_epochs', [], 'Specify a list of training epochs, pause training after ' 'each of these epochs, run evaluation, then resume training. ' 'See --eval_during_training_every_n_steps for more ' 'information.') flags.DEFINE_boolean('forward_only', False, 'whether use forward-only or training for benchmarking') flags.DEFINE_boolean('freeze_when_forward_only', False, 'whether to freeze the graph when in forward-only mode.') flags.DEFINE_boolean('print_training_accuracy', False, 'whether to calculate and print training accuracy during ' 'training') flags.DEFINE_integer('batch_size', 0, 'batch size per compute device') flags.DEFINE_integer('eval_batch_size', 0, 'eval batch size per compute device') flags.DEFINE_integer('batch_group_size', 1, 'number of groups of batches processed in the image ' 'producer.') flags.DEFINE_integer('num_batches', None, 'number of batches to run, excluding ' 'warmup. Defaults to %d' % _DEFAULT_NUM_BATCHES) flags.DEFINE_integer('num_eval_batches', None, 'number of eval batches to run, excluding warmup. ' 'Defaults to --num_batches') flags.DEFINE_float('num_epochs', None, 'number of epochs to run, excluding warmup. ' 'This and --num_batches cannot both be specified.') flags.DEFINE_float('num_eval_epochs', None, 'number of eval epochs to run, excluding warmup. ' 'Defaults to --num_epochs') flags.DEFINE_float('stop_at_top_1_accuracy', None, 'If set, stops training after the evaluation accuracy hits ' 'this number. Can only be used with one of the ' '--eval_during_training_* flags.') flags.DEFINE_boolean('collect_eval_results_async', False, 'If True, start a separate process to postprocess eval ' 'results asynchronously. This currently only works with ' 'the SSD model.') flags.DEFINE_integer('num_warmup_batches', None, 'number of batches to run before timing') flags.DEFINE_integer('autotune_threshold', None, 'The autotune threshold for the models') # TODO(tucker): change num_gpus to num_devices flags.DEFINE_integer('num_gpus', 1, 'the number of GPUs to run on') flags.DEFINE_string('gpu_indices', '', 'indices of worker GPUs in ring order') flags.DEFINE_integer('display_every', 10, 'Number of local steps after which progress is printed ' 'out') flags.DEFINE_float('display_perf_ewma', None, 'If set, display numbers of images/sec using exponentially ' 'weighted moving avearge with the specified weight, which ' 'defines how much current value contributes to the reported ' 'average. Increasing weight makes the reported performance ' 'number reflect more about the real-time speed instead of ' 'the entire history', lower_bound=0, upper_bound=1) flags.DEFINE_string('data_dir', None, 'Path to dataset in TFRecord format (aka Example ' 'protobufs). If not specified, synthetic data will be ' 'used.') flags.DEFINE_string('data_name', None, 'Name of dataset: imagenet or cifar10. If not specified, ' 'it is automatically guessed based on data_dir.') flags.DEFINE_string('resize_method', 'bilinear', 'Method for resizing input images: crop, nearest, ' 'bilinear, bicubic, area, or round_robin. The `crop` mode ' 'requires source images to be at least as large as the ' 'network input size. The `round_robin` mode applies ' 'different resize methods based on position in a batch in ' 'a round-robin fashion. Other modes support any sizes and ' 'apply random bbox distortions before resizing (even with ' 'distortions=False).') flags.DEFINE_boolean('distortions', False, 'Enable/disable distortions during image preprocessing. ' 'These include bbox and color distortions.') flags.DEFINE_boolean('use_datasets', True, 'Enable use of datasets for input pipeline') flags.DEFINE_string('input_preprocessor', 'default', 'Name of input preprocessor. The list of supported input ' 'preprocessors are defined in preprocessing.py.') flags.DEFINE_string('gpu_thread_mode', 'gpu_private', 'Methods to assign GPU host work to threads. ' 'global: all GPUs and CPUs share the same global threads; ' 'gpu_private: a private threadpool for each GPU; ' 'gpu_shared: all GPUs share the same threadpool.') flags.DEFINE_integer('per_gpu_thread_count', 0, 'The number of threads to use for GPU. Only valid when ' 'gpu_thread_mode is not global.') flags.DEFINE_boolean('hierarchical_copy', False, 'Use hierarchical copies. Currently only optimized for ' 'use on a DGX-1 with 8 GPUs and may perform poorly on ' 'other hardware. Requires --num_gpus > 1, and only ' 'recommended when --num_gpus=8') # TODO(hinsu): Support auto-detection of the network topology while still # retaining the ability to specify a particular topology for debugging. flags.DEFINE_enum( 'network_topology', constants.NetworkTopology.DGX1, (constants.NetworkTopology.DGX1, constants.NetworkTopology.GCP_V100), 'Network topology specifies the topology used to connect multiple devices. ' 'Network topology is used to decide the hierarchy to use for the ' 'hierarchical_copy.') flags.DEFINE_integer('gradient_repacking', 0, 'Use gradient repacking. It ' 'currently only works with replicated mode. At the end of ' 'each step, it repacks the gradients for more efficient ' 'cross-device transportation. A non-zero value specifies ' 'the number of split packs that will be formed.', lower_bound=0) flags.DEFINE_boolean('compact_gradient_transfer', True, 'Compact gradient' 'as much as possible for cross-device transfer and ' 'aggregation.') flags.DEFINE_enum('variable_consistency', 'strong', ('strong', 'relaxed'), 'The data consistency for trainable variables. With strong ' 'consistency, the variable always have the updates from ' 'previous step. With relaxed consistency, all the updates ' 'will eventually show up in the variables. Likely one step ' 'behind.') flags.DEFINE_boolean('datasets_repeat_cached_sample', False, 'Enable use of a special datasets pipeline that reads a ' 'single TFRecord into memory and repeats it infinitely ' 'many times. The purpose of this flag is to make it ' 'possible to write regression tests that are not ' 'bottlenecked by CNS throughput. ' 'Use datasets_use_caching to cache input data.') flags.DEFINE_enum('local_parameter_device', 'gpu', ('cpu', 'gpu', 'CPU', 'GPU'), 'Device to use as parameter server: cpu or gpu. For ' 'distributed training, it can affect where caching of ' 'variables happens.') flags.DEFINE_enum('device', 'gpu', ('cpu', 'gpu', 'CPU', 'GPU'), 'Device to use for computation: cpu or gpu') flags.DEFINE_enum('data_format', 'NCHW', ('NHWC', 'NCHW'), 'Data layout to use: NHWC (TF native) or NCHW (cuDNN ' 'native, requires GPU).') flags.DEFINE_integer('num_intra_threads', None, 'Number of threads to use for intra-op parallelism. If ' 'set to 0, the system will pick an appropriate number. ' 'None is the same as 0 except that it disables intra-op ' 'parallelism on a GPU.') flags.DEFINE_integer('num_inter_threads', 0, 'Number of threads to use for inter-op parallelism. If ' 'set to 0, the system will pick an appropriate number.') flags.DEFINE_boolean('use_numa_affinity', False, 'Whether to turn on NUMA affinity for CPU devices. ' 'This is probably only useful when --device=cpu.') flags.DEFINE_string('trace_file', '', 'Enable TensorFlow tracing and write trace to this file.') flags.DEFINE_boolean('use_chrome_trace_format', True, 'If True, the trace_file, if specified, will be in a ' 'Chrome trace format. If False, then it will be a ' 'StepStats raw proto.') _NUM_STEPS_TO_PROFILE = 10 _NUM_OPS_TO_PRINT = 20 flags.DEFINE_string('tfprof_file', None, 'If specified, write a tfprof ProfileProto to this file. ' 'The performance and other aspects of the model can then ' 'be analyzed with tfprof. See ' 'https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/g3doc/command_line.md ' # pylint: disable=line-too-long 'for more info on how to do this. The first %d steps ' 'are profiled. Additionally, the top %d most time ' 'consuming ops will be printed.\n' 'Note: profiling with tfprof is very slow, but most of the ' 'overhead is spent between steps. So, profiling results ' 'are more accurate than the slowdown would suggest.' % (_NUM_STEPS_TO_PROFILE, _NUM_OPS_TO_PRINT)) flags.DEFINE_string('graph_file', None, 'Write the model\'s graph definition to this file. ' 'Defaults to binary format unless filename ends in "txt".') flags.DEFINE_string('partitioned_graph_file_prefix', None, 'If specified, after the graph has been partitioned and ' 'optimized, write out each partitioned graph to a file ' 'with the given prefix.') flags.DEFINE_enum('optimizer', 'sgd', ('momentum', 'sgd', 'rmsprop', 'adam'), 'Optimizer to use') flags.DEFINE_float('init_learning_rate', None, 'Initial learning rate for training.') flags.DEFINE_string('piecewise_learning_rate_schedule', None, 'Specifies a piecewise learning rate schedule based on the ' 'number of epochs. This is the form LR0;E1;LR1;...;En;LRn, ' 'where each LRi is a learning rate and each Ei is an epoch ' 'indexed from 0. The learning rate is LRi if the ' 'E(i-1) <= current_epoch < Ei. For example, if this ' 'paramater is 0.3;10;0.2;25;0.1, the learning rate is 0.3 ' 'for the first 10 epochs, then is 0.2 for the next 15 ' 'epochs, then is 0.1 until training ends.') flags.DEFINE_float('num_epochs_per_decay', 0, 'Steps after which learning rate decays. If 0, the learning ' 'rate does not decay.') flags.DEFINE_float('learning_rate_decay_factor', 0, 'Learning rate decay factor. Decay by this factor every ' '`num_epochs_per_decay` epochs. If 0, learning rate does ' 'not decay.') flags.DEFINE_float('num_learning_rate_warmup_epochs', 0, 'Slowly increase to the initial learning rate in the first ' 'num_learning_rate_warmup_epochs linearly.') flags.DEFINE_float('minimum_learning_rate', 0, 'The minimum learning rate. The learning rate will ' 'never decay past this value. Requires `learning_rate`, ' '`num_epochs_per_decay` and `learning_rate_decay_factor` to ' 'be set.') flags.DEFINE_float('resnet_base_lr', None, "Base learning rate at bs=256. Only " "relevant when training ResNet and utilizing the model's " "learning rate heuristic (get_learning_rate).") flags.DEFINE_float('momentum', 0.9, 'Momentum for training.') flags.DEFINE_float('rmsprop_decay', 0.9, 'Decay term for RMSProp.') flags.DEFINE_float('rmsprop_momentum', 0.9, 'Momentum in RMSProp.') flags.DEFINE_float('rmsprop_epsilon', 1.0, 'Epsilon term for RMSProp.') flags.DEFINE_float('adam_beta1', 0.9, 'Beta2 term for the Adam optimizer') flags.DEFINE_float('adam_beta2', 0.999, 'Beta2 term for the Adam optimizer') flags.DEFINE_float('adam_epsilon', 1e-8, 'Epsilon term for the Adam optimizer') flags.DEFINE_float('gradient_clip', None, 'Gradient clipping magnitude. Disabled by default.') flags.DEFINE_float('weight_decay', 0.00004, 'Weight decay factor for training.') flags.DEFINE_float('gpu_memory_frac_for_testing', 0, 'If non-zero, the fraction of GPU memory that will be used. ' 'Useful for testing the benchmark script, as this allows ' 'distributed mode to be run on a single machine. For ' 'example, if there are two tasks, each can be allocated ' '~40 percent of the memory on a single machine. This is ' 'also useful for using unified memory, as this can be set ' 'above 1 to oversubscribe the GPU using unified memory.', lower_bound=0.) flags.DEFINE_boolean('use_unified_memory', None, 'If True, allocate unified memory enabling larger models ' 'to fit in available device RAM.') flags.DEFINE_boolean('timestamped_allocator', False, 'If True marks free BFCAllocator::Chunks with time ' 'at which they are freed which can allow more efficient ' 'memory allocation in cases like RDMA networking.') flags.DEFINE_integer('gpu_kt_max_interval', 0, 'If > 0, the maximum number of GPU Ops that may be queued ' 'in a row without also queuing a tracking event.') flags.DEFINE_integer('gpu_kt_max_bytes', 0, 'If > 0, the maximum number of bytes ' 'of GPU memory that may be allocated by sequential ' 'GPU Ops without queuing a tracking event.') flags.DEFINE_integer('gpu_kt_max_pending', 0, 'If > 0 no more than this many GPU tracking events may be ' 'outstanding at any time. When this limit is reached ' 'launch of additional kernels will stall until an ' 'outstanding event completes.') flags.DEFINE_boolean('use_tf_layers', True, 'If True, use tf.layers for neural network layers. This ' 'should not affect performance or accuracy in any way.') flags.DEFINE_integer('tf_random_seed', 1234, 'The TensorFlow random seed. Useful for debugging NaNs, ' 'as this can be set to various values to see if the NaNs ' 'depend on the seed.') flags.DEFINE_string('debugger', None, 'If set, use the TensorFlow debugger. If set to "cli", use ' 'the local CLI debugger. Otherwise, this must be in the ' 'form hostname:port (e.g., localhost:7007) in which case ' 'the experimental TensorBoard debugger will be used') flags.DEFINE_boolean('use_python32_barrier', False, 'When on, use threading.Barrier at Python 3.2.') flags.DEFINE_boolean('ml_perf', False, 'When True, change how the Imagenet input pipeline works ' 'slightly to meet the MLPerf compliance rules. This slows ' 'down the input pipeline. Without this option, at the end ' 'of the input pipeline, the image is divided by 127.5, ' 'then 1.0 is subtracted from it, bringing the image ' 'values from [0, 255] to [-1.0, 1.0]. With this option, ' 'each of the three channels (red, green, blue) have the ' 'average channel value among all image subtracted from ' 'it, and no division is done.') flags.DEFINE_boolean('datasets_use_prefetch', True, 'Enable use of prefetched datasets for input pipeline. ' 'This option is meaningless if use_datasets=False.') flags.DEFINE_integer('datasets_prefetch_buffer_size', 1, 'Prefetching op buffer size per compute device.') flags.DEFINE_integer('datasets_num_private_threads', None, 'Number of threads for a private threadpool created for ' 'all datasets computation. By default, we pick an ' 'appropriate number. If set to 0, we use the default ' 'tf-Compute threads for dataset operations.') flags.DEFINE_boolean('datasets_use_caching', False, 'Cache the compressed input data in memory. This improves ' 'the data input performance, at the cost of additional ' 'memory.') flags.DEFINE_integer('datasets_parallel_interleave_cycle_length', None, 'Number of parallel file readers interleaving input data.') flags.DEFINE_boolean('datasets_sloppy_parallel_interleave', False, 'Allow parallel interleave to depart from deterministic ' 'ordering, by temporarily skipping over files whose ' 'elements are not readily available. This can increase ' 'througput in particular in the presence of stragglers.') flags.DEFINE_integer('datasets_parallel_interleave_prefetch', None, 'The number of input elements to fetch before they are ' 'needed for interleaving.') flags.DEFINE_integer( 'multi_device_iterator_max_buffer_size', 1, 'Configuration parameter for the MultiDeviceIterator that ' ' specifies the host side buffer size for each device.') # Performance tuning parameters. flags.DEFINE_boolean('winograd_nonfused', True, 'Enable/disable using the Winograd non-fused algorithms.') flags.DEFINE_boolean( 'batchnorm_persistent', True, 'Enable/disable using the CUDNN_BATCHNORM_SPATIAL_PERSISTENT ' 'mode for batchnorm.') flags.DEFINE_boolean('sync_on_finish', False, 'Enable/disable whether the devices are synced after each ' 'step.') flags.DEFINE_boolean('staged_vars', False, 'whether the variables are staged from the main ' 'computation') flags.DEFINE_boolean('force_gpu_compatible', False, 'whether to enable force_gpu_compatible in GPU_Options') flags.DEFINE_boolean('allow_growth', None, 'whether to enable allow_growth in GPU_Options') flags.DEFINE_boolean('xla', False, 'whether to enable XLA auto-jit compilation') flags.DEFINE_boolean('xla_compile', False, 'Enable xla to compile the graph. Uncompilable ops will ' 'result in fatal errors.') flags.DEFINE_boolean('fuse_decode_and_crop', True, 'Fuse decode_and_crop for image preprocessing.') flags.DEFINE_boolean('distort_color_in_yiq', True, 'Distort color of input images in YIQ space.') flags.DEFINE_boolean('enable_optimizations', True, 'Whether to enable grappler and other optimizations.') flags.DEFINE_string('rewriter_config', None, 'Config for graph optimizers, described as a ' 'RewriterConfig proto buffer.') flags.DEFINE_enum('loss_type_to_report', 'total_loss', ('base_loss', 'total_loss'), 'Which type of loss to output and to write summaries for. ' 'The total loss includes L2 loss while the base loss does ' 'not. Note that the total loss is always used while ' 'computing gradients during training if weight_decay > 0, ' 'but explicitly computing the total loss, instead of just ' 'computing its gradients, can have a performance impact.') flags.DEFINE_boolean('single_l2_loss_op', False, 'If True, instead of using an L2 loss op per variable, ' 'concatenate the variables into a single tensor and do a ' 'single L2 loss on the concatenated tensor.') flags.DEFINE_boolean('use_resource_vars', False, 'Use resource variables instead of normal variables. ' 'Resource variables are slower, but this option is useful ' 'for debugging their performance.') flags.DEFINE_boolean('compute_lr_on_cpu', False, 'If True, do computations related to learning rate on the ' 'CPU instead of the GPU. This will significantly improve ' 'XLA performance in some cases.') flags.DEFINE_boolean('sparse_to_dense_grads', False, 'If True, convert all sparse gradients to dense gradients ' 'before passing them to the optimizer to update ' 'variables. Only affects models with sparse gradients, ' 'which currently is only the NCF model.') # Performance tuning specific to MKL. flags.DEFINE_boolean('mkl', False, 'If true, set MKL environment variables.') flags.DEFINE_integer('kmp_blocktime', 0, 'The time, in milliseconds, that a thread should wait, ' 'after completing the execution of a parallel region, ' 'before sleeping') flags.DEFINE_string('kmp_affinity', 'granularity=fine,verbose,compact,1,0', 'Restricts execution of certain threads (virtual execution ' 'units) to a subset of the physical processing units in a ' 'multiprocessor computer.') flags.DEFINE_integer('kmp_settings', 1, 'If set to 1, MKL settings will be printed.') # fp16 parameters. If use_fp16=False, no other fp16 parameters apply. flags.DEFINE_boolean('use_fp16', False, 'Use 16-bit floats for certain tensors instead of 32-bit ' 'floats. This is currently experimental.') # TODO(reedwm): The default loss scale of 128 causes most models to diverge # on the second step with synthetic data. Changing the tf.set_random_seed # call to tf.set_random_seed(1235) or most other seed values causes the # issue not to occur. flags.DEFINE_float('fp16_loss_scale', None, 'If fp16 is enabled, the loss is multiplied by this amount ' 'right before gradients are computed, then each gradient ' 'is divided by this amount. Mathematically, this has no ' 'effect, but it helps avoid fp16 underflow. Set to 1 to ' 'effectively disable. Ignored during eval.') flags.DEFINE_boolean('fp16_vars', False, 'If fp16 is enabled, also use fp16 for variables. If ' 'False, the variables are stored in fp32 and casted to ' 'fp16 when retrieved. Recommended to leave as False.') flags.DEFINE_boolean('fp16_enable_auto_loss_scale', False, 'If True and use_fp16 is True, automatically adjust the ' 'loss scale during training.') flags.DEFINE_integer('fp16_inc_loss_scale_every_n', 1000, 'If fp16 is enabled and fp16_enable_auto_loss_scale is ' 'True, increase the loss scale every n steps.') # The method for managing variables: # parameter_server: variables are stored on a parameter server that holds # the master copy of the variable. In local execution, a local device # acts as the parameter server for each variable; in distributed # execution, the parameter servers are separate processes in the # cluster. # For each step, each tower gets a copy of the variables from the # parameter server, and sends its gradients to the param server. # replicated: each GPU has its own copy of the variables. To apply # gradients, an all_reduce algorithm or or regular cross-device # aggregation is used to replicate the combined gradients to all # towers (depending on all_reduce_spec parameter setting). # independent: each GPU has its own copy of the variables, and gradients # are not shared between towers. This can be used to check performance # when no data is moved between GPUs. # distributed_replicated: Distributed training only. Each GPU has a copy # of the variables, and updates its copy after the parameter servers # are all updated with the gradients from all servers. Only works with # cross_replica_sync=true. Unlike 'replicated', currently never uses # nccl all-reduce for replicating within a server. # distributed_all_reduce: Distributed training where all replicas run # in a single session, using all-reduce to mutally reduce the # gradients. Uses no parameter servers. When there is only one # worker, this is the same as replicated. # collective_all_reduce: Distributed training where all replicas run # independepently except for variable initialization and for # gradient reduction which is done via collective all-reduce. # NOTE: collective_all_reduce in conjunction with use_fp16 can # lead to NaNs in some models (resnet50). TODO(tucker): fix it. # horovod: Distributed training using Horovod library. Runs workers using # an MPI framework (e.g. Open MPI). Each worker runs training on # single GPU, and averages gradients using NCCL or MPI all-reduce. # See https://github.com/uber/horovod for more details. flags.DEFINE_enum('variable_update', 'parameter_server', ('parameter_server', 'replicated', 'distributed_replicated', 'independent', 'distributed_all_reduce', 'collective_all_reduce', 'horovod'), 'The method for managing variables: parameter_server, ' 'replicated, distributed_replicated, independent, ' 'distributed_all_reduce, collective_all_reduce, horovod') flags.DEFINE_string('all_reduce_spec', None, 'A specification of the all_reduce algorithm to be used ' 'for reducing gradients. For more details, see ' 'parse_all_reduce_spec in variable_mgr.py. An ' 'all_reduce_spec has BNF form:\n' 'int ::= positive whole number\n' 'g_int ::= int[KkMGT]?\n' 'alg_spec ::= alg | alg#int\n' 'range_spec ::= alg_spec | alg_spec/alg_spec\n' 'spec ::= range_spec | range_spec:g_int:range_spec\n' 'NOTE: not all syntactically correct constructs are ' 'supported.\n\n' 'Examples:\n ' '"xring" == use one global ring reduction for all ' 'tensors\n' '"pscpu" == use CPU at worker 0 to reduce all tensors\n' '"nccl" == use NCCL to locally reduce all tensors. ' 'Limited to 1 worker.\n' '"nccl/xring" == locally (to one worker) reduce values ' 'using NCCL then ring reduce across workers.\n' '"pscpu:32k:xring" == use pscpu algorithm for tensors of ' 'size up to 32kB, then xring for larger tensors.') # If variable_update==distributed_all_reduce then it may be advantageous # to aggregate small tensors into one prior to reduction. These parameters # control that aggregation. flags.DEFINE_integer('agg_small_grads_max_bytes', 0, 'If > 0, try to aggregate tensors of less than this ' 'number of bytes prior to all-reduce.') flags.DEFINE_integer('agg_small_grads_max_group', 10, 'When aggregating small tensors for all-reduce do not ' 'aggregate more than this many into one new tensor.') flags.DEFINE_integer('allreduce_merge_scope', 1, 'Establish a name scope around this many ' 'gradients prior to creating the all-reduce operations. ' 'It may affect the ability of the backend to merge ' 'parallel ops.') # Distributed training parameters. flags.DEFINE_enum('job_name', '', ('ps', 'worker', 'controller', ''), 'One of "ps", "worker", "controller", "". Empty for local ' 'training') flags.DEFINE_string('ps_hosts', '', 'Comma-separated list of target hosts') flags.DEFINE_string('worker_hosts', '', 'Comma-separated list of target hosts') flags.DEFINE_string('controller_host', None, 'optional controller host') flags.DEFINE_integer('task_index', 0, 'Index of task within the job') flags.DEFINE_string('server_protocol', 'grpc', 'protocol for servers') flags.DEFINE_boolean('cross_replica_sync', True, '') flags.DEFINE_string('horovod_device', '', 'Device to do Horovod all-reduce on: ' 'empty (default), cpu or gpu. Default with utilize GPU if ' 'Horovod was compiled with the HOROVOD_GPU_ALLREDUCE ' 'option, and CPU otherwise.') # Summary and Save & load checkpoints. flags.DEFINE_integer('summary_verbosity', 0, 'Verbosity level for summary ops. ' 'level 0: disable any summary.\n' 'level 1: small and fast ops, e.g.: learning_rate, ' 'total_loss.\n' 'level 2: medium-cost ops, e.g. histogram of all ' 'gradients.\n' 'level 3: expensive ops: images and histogram of each ' 'gradient.\n') flags.DEFINE_integer('save_summaries_steps', 0, 'How often to save summaries for trained models. Pass 0 ' 'to disable summaries.') flags.DEFINE_integer('save_model_secs', 0, 'How often to save trained models. Pass 0 to disable ' 'saving checkpoints every N seconds. A checkpoint is ' 'saved after training completes regardless of this ' 'option.') flags.DEFINE_integer('save_model_steps', None, 'How often to save trained models. If specified, ' 'save_model_secs must not be specified.') flags.DEFINE_integer('max_ckpts_to_keep', 5, 'Max number of checkpoints to keep.') flags.DEFINE_string('train_dir', None, 'Path to session checkpoints. Pass None to disable saving ' 'checkpoint at the end.') flags.DEFINE_string('eval_dir', '/tmp/tf_cnn_benchmarks/eval', 'Directory where to write eval event logs.') flags.DEFINE_string('backbone_model_path', None, 'Path to pretrained backbone model checkpoint. Pass None ' 'if not using a backbone model.') flags.DEFINE_enum('trt_mode', '', ['', 'FP32', 'FP16', 'INT8'], 'If this is specified in forward_only mode and ' 'freeze_when_forward_only is set to True, use TensorRT to ' 'optimize the graph before execution.') flags.DEFINE_integer('trt_max_workspace_size_bytes', 4 << 30, 'Max workspace size bytes used by the TensorRT optimizer.') # Benchmark logging for model garden metric flags.DEFINE_string('benchmark_log_dir', None, 'The directory to place the log files containing the ' 'results of benchmark. The logs are created by ' 'BenchmarkFileLogger. Requires the root of the Tensorflow ' 'models repository to be in $PYTHTONPATH.') flags.DEFINE_string('benchmark_test_id', None, 'The unique test ID of the benchmark run. It could be the ' 'combination of key parameters. It is hardware independent ' 'and could be used compare the performance between ' 'different test runs. This flag is designed for human ' 'consumption, and does not have any impact within the ' 'system.') platforms_util.define_platform_params() class GlobalStepWatcher(threading.Thread): """A helper class for global_step. Polls for changes in the global_step of the model, and finishes when the number of steps for the global run are done. """ def __init__(self, sess, global_step_op, start_at_global_step, end_at_global_step): threading.Thread.__init__(self) self.sess = sess self.global_step_op = global_step_op self.start_at_global_step = start_at_global_step self.end_at_global_step = end_at_global_step self.start_time = 0 self.start_step = 0 self.finish_time = 0 self.finish_step = 0 def run(self): while self.finish_time == 0: time.sleep(.25) global_step_val, = self.sess.run([self.global_step_op]) if self.start_time == 0 and global_step_val >= self.start_at_global_step: # Use tf.logging.info instead of log_fn, since print (which is log_fn) # is not thread safe and may interleave the outputs from two parallel # calls to print, which can break tests. tf.logging.info('Starting real work at step %s at time %s' % (global_step_val, time.ctime())) self.start_time = time.perf_counter() self.start_step = global_step_val if self.finish_time == 0 and global_step_val >= self.end_at_global_step: tf.logging.info('Finishing real work at step %s at time %s' % (global_step_val, time.ctime())) self.finish_time = time.perf_counter() self.finish_step = global_step_val def done(self): return self.finish_time > 0 def num_steps(self): return self.finish_step - self.start_step def elapsed_time(self): return self.finish_time - self.start_time class CheckpointNotFoundException(Exception): pass def create_config_proto(params): """Returns session config proto. Args: params: Params tuple, typically created by make_params or make_params_from_flags. """ config = tf.ConfigProto() config.allow_soft_placement = True if params.num_intra_threads is None: if params.device == 'gpu': config.intra_op_parallelism_threads = 1 else: config.intra_op_parallelism_threads = params.num_intra_threads config.inter_op_parallelism_threads = params.num_inter_threads config.experimental.collective_group_leader = '/job:worker/replica:0/task:0' config.gpu_options.experimental.collective_ring_order = params.gpu_indices config.gpu_options.force_gpu_compatible = params.force_gpu_compatible config.experimental.use_numa_affinity = params.use_numa_affinity if params.device == 'cpu': # TODO(tucker): change num_gpus to num_devices config.device_count['CPU'] = params.num_gpus if params.allow_growth is not None: config.gpu_options.allow_growth = params.allow_growth if params.gpu_memory_frac_for_testing > 0: config.gpu_options.per_process_gpu_memory_fraction = ( params.gpu_memory_frac_for_testing) if params.use_unified_memory: config.gpu_options.experimental.use_unified_memory = ( params.use_unified_memory) if params.timestamped_allocator: config.gpu_options.experimental.timestamped_allocator = ( params.timestamped_allocator) if params.gpu_kt_max_interval > 0: config.gpu_options.experimental.kernel_tracker_max_interval = ( params.gpu_kt_max_interval) if params.gpu_kt_max_bytes > 0: config.gpu_options.experimental.kernel_tracker_max_bytes = ( params.gpu_kt_max_bytes) if params.gpu_kt_max_pending > 0: config.gpu_options.experimental.kernel_tracker_max_pending = ( params.gpu_kt_max_pending) if params.xla: config.graph_options.optimizer_options.global_jit_level = ( tf.OptimizerOptions.ON_1) if params.rewriter_config: rewriter_config = rewriter_config_pb2.RewriterConfig() text_format.Merge(params.rewriter_config, rewriter_config) config.graph_options.rewrite_options.CopyFrom(rewriter_config) elif not params.enable_optimizations: config.graph_options.optimizer_options.opt_level = tf.OptimizerOptions.L0 config.graph_options.rewrite_options.disable_meta_optimizer = True elif params.variable_update == 'collective_all_reduce': rewrite_options = config.graph_options.rewrite_options rewrite_options.scoped_allocator_optimization = ( rewriter_config_pb2.RewriterConfig.ON) rewrite_options.scoped_allocator_opts.enable_op.append('CollectiveReduce') if params.variable_update == 'horovod': import horovod.tensorflow as hvd # pylint: disable=g-import-not-at-top config.gpu_options.visible_device_list = str(hvd.local_rank()) # For collective_all_reduce, ignore all devices except current worker. if params.variable_update == 'collective_all_reduce': del config.device_filters[:] config.device_filters.append( '/job:%s/replica:0/task:%d' % (params.job_name, params.task_index)) # TODO(b/117324590): Re-enable PinToHostOptimizer when b/117324590 is fixed. # Currently we have to disable PinToHostOptimizer w/ XLA since it causes # OOM/perf cliffs. config.graph_options.rewrite_options.pin_to_host_optimization = ( rewriter_config_pb2.RewriterConfig.OFF) return config def get_mode_from_params(params): """Returns the mode in which this script is running. Args: params: Params tuple, typically created by make_params or make_params_from_flags. Raises: ValueError: Unsupported params settings. """ if params.forward_only and params.eval: raise ValueError('Only one of forward_only and eval parameters is true') if params.eval: return constants.BenchmarkMode.EVAL elif params.forward_only: return constants.BenchmarkMode.FORWARD_ONLY elif (params.eval_during_training_every_n_steps or params.eval_during_training_every_n_epochs or params.eval_during_training_at_specified_steps or params.eval_during_training_at_specified_epochs): return constants.BenchmarkMode.TRAIN_AND_EVAL else: return constants.BenchmarkMode.TRAIN # How many digits to show for the loss and accuracies during training. LOSS_AND_ACCURACY_DIGITS_TO_SHOW = 3 def benchmark_one_step(sess, fetches, step, batch_size, step_train_times, trace_filename, partitioned_graph_file_prefix, profiler, image_producer, params, summary_op=None, show_images_per_sec=True, benchmark_logger=None, collective_graph_key=0, should_output_files=True): """Advance one step of benchmarking.""" should_profile = profiler and 0 <= step < _NUM_STEPS_TO_PROFILE need_options_and_metadata = ( should_profile or collective_graph_key > 0 or ((trace_filename or partitioned_graph_file_prefix) and step == -2) ) if need_options_and_metadata: run_options = tf.RunOptions() if (trace_filename and step == -2) or should_profile: run_options.trace_level = tf.RunOptions.FULL_TRACE if partitioned_graph_file_prefix and step == -2: run_options.output_partition_graphs = True if collective_graph_key > 0: run_options.experimental.collective_graph_key = collective_graph_key run_metadata = tf.RunMetadata() else: run_options = None run_metadata = None summary_str = None start_time = time.perf_counter() if summary_op is None: results = sess.run(fetches, options=run_options, run_metadata=run_metadata) else: (results, summary_str) = sess.run( [fetches, summary_op], options=run_options, run_metadata=run_metadata) if not params.forward_only: lossval = results['average_loss'] else: lossval = 0. if image_producer is not None: image_producer.notify_image_consumption() train_time = time.perf_counter() - start_time step_train_times.append(train_time) if (show_images_per_sec and step >= 0 and (step == 0 or (step + 1) % params.display_every == 0)): speed_mean, speed_uncertainty, speed_jitter = get_perf_timing( batch_size, step_train_times, params.display_perf_ewma) log_str = '%i\t%s\t%.*f' % ( step + 1, get_perf_timing_str(speed_mean, speed_uncertainty, speed_jitter), LOSS_AND_ACCURACY_DIGITS_TO_SHOW, lossval) if 'top_1_accuracy' in results: log_str += '\t%.*f\t%.*f' % ( LOSS_AND_ACCURACY_DIGITS_TO_SHOW, results['top_1_accuracy'], LOSS_AND_ACCURACY_DIGITS_TO_SHOW, results['top_5_accuracy']) log_fn(log_str) if benchmark_logger: benchmark_logger.log_metric( 'current_examples_per_sec', speed_mean, global_step=step + 1) if 'top_1_accuracy' in results: benchmark_logger.log_metric( 'top_1_accuracy', results['top_1_accuracy'], global_step=step + 1) benchmark_logger.log_metric( 'top_5_accuracy', results['top_5_accuracy'], global_step=step + 1) if need_options_and_metadata: if should_profile: profiler.add_step(step, run_metadata) if trace_filename and step == -2 and should_output_files: log_fn('Dumping trace to %s' % trace_filename) trace_dir = os.path.dirname(trace_filename) if not gfile.Exists(trace_dir): gfile.MakeDirs(trace_dir) with gfile.Open(trace_filename, 'w') as trace_file: if params.use_chrome_trace_format: trace = timeline.Timeline(step_stats=run_metadata.step_stats) trace_file.write(trace.generate_chrome_trace_format(show_memory=True)) else: trace_file.write(str(run_metadata.step_stats)) if partitioned_graph_file_prefix and step == -2 and should_output_files: path, filename = os.path.split(partitioned_graph_file_prefix) if '.' in filename: base_filename, ext = filename.rsplit('.', 1) ext = '.' + ext else: base_filename, ext = filename, '' as_text = filename.endswith('txt') for graph_def in run_metadata.partition_graphs: device = graph_def.node[0].device.replace('/', '_').replace(':', '_') graph_filename = '%s%s%s' % (base_filename, device, ext) log_fn('Writing partitioned GraphDef as %s to %s' % ( 'text' if as_text else 'binary', os.path.join(path, graph_filename))) tf.train.write_graph(graph_def, path, graph_filename, as_text) return (summary_str, lossval) def get_perf_timing_str(speed_mean, speed_uncertainty, speed_jitter, scale=1): if scale == 1: # TODO(laigd): rename 'images' to maybe 'inputs', same below. return ('images/sec: %.1f +/- %.1f (jitter = %.1f)' % (speed_mean, speed_uncertainty, speed_jitter)) else: return 'images/sec: %.1f' % speed_mean def get_perf_timing(batch_size, step_train_times, ewma_alpha=None, scale=1): """Calculate benchmark processing speed.""" times = np.array(step_train_times) speeds = batch_size / times if ewma_alpha: weights = np.logspace(len(times)-1, 0, len(times), base=1-ewma_alpha) time_mean = np.average(times, weights=weights) else: time_mean = np.mean(times) speed_mean = scale * batch_size / time_mean speed_uncertainty = np.std(speeds) / np.sqrt(float(len(speeds))) speed_jitter = 1.4826 * np.median(np.abs(speeds - np.median(speeds))) return speed_mean, speed_uncertainty, speed_jitter def load_checkpoint(saver, sess, ckpt_dir): """Loads checkpoint from provided directory or full path. Args: saver: Saver used to restore the checkpoint. sess: TensorFlow session. ckpt_dir: Path to a folder of checkpoints or full path to a checkpoint. Returns: Global step. """ model_checkpoint_path = _get_checkpoint_to_load(ckpt_dir) global_step = model_checkpoint_path.split('/')[-1].split('-')[-1] if not global_step.isdigit(): global_step = 0 else: global_step = int(global_step) saver.restore(sess, model_checkpoint_path) log_fn('Successfully loaded model from %s.' % model_checkpoint_path) return global_step def _get_checkpoint_to_load(ckpt_dir): """Returns which checkpoint to load. Args: ckpt_dir: Path to a folder of checkpoints or full path to a checkpoint. Returns: Full path to checkpoint to load. Raises: CheckpointNotFoundException: If checkpoint is not found. """ p = re.compile(r'ckpt-\d+$') if p.search(ckpt_dir): model_checkpoint_path = ckpt_dir else: # Finds latest checkpoint in directory provided ckpt = tf.train.get_checkpoint_state(ckpt_dir) if ckpt and ckpt.model_checkpoint_path: model_checkpoint_path = ckpt.model_checkpoint_path else: raise CheckpointNotFoundException('No checkpoint file found in dir:{}'. format(ckpt_dir)) return model_checkpoint_path # Params are passed to BenchmarkCNN's constructor. Params is a map from name # to value, with one field per key in flags.param_specs. # # Call make_params() or make_params_from_flags() below to construct a Params # tuple with default values from flags.param_specs, rather than constructing # Params directly. Params = namedtuple('Params', flags.param_specs.keys()) # pylint: disable=invalid-name def validate_params(params): """Validates that the Params tuple had valid values. When command-line flags are defined for each ParamSpec by calling flags.define_flags(), calling this function is unnecessary because absl already does flag validation. Otherwise, this function should be called. Args: params: A Params tuple. Raises: ValueError: An element of params had an invalid value. """ for name, value in params._asdict().items(): param_spec = flags.param_specs[name] if param_spec.flag_type in ('integer', 'float'): if (value is not None and param_spec.kwargs['lower_bound'] is not None and value < param_spec.kwargs['lower_bound']): raise ValueError('Param %s value of %s is lower than the lower bound ' 'of %s' % (name, value, param_spec.kwargs['lower_bound'])) if (value is not None and param_spec.kwargs['upper_bound'] is not None and param_spec.kwargs['upper_bound'] < value): raise ValueError('Param %s value of %s is higher than the upper bound ' 'of %s' % (name, value, param_spec.kwargs['upper_bound'])) elif (value is not None and param_spec.flag_type == 'enum' and value not in param_spec.kwargs['enum_values']): raise ValueError('Param %s of value %s is not in %s'% (name, value, param_spec.kwargs['enum_values'])) def make_params(**kwargs): """Create a Params tuple for BenchmarkCNN from kwargs. Default values are filled in from flags.param_specs. Args: **kwargs: kwarg values will override the default values. Returns: Params namedtuple for constructing BenchmarkCNN. """ # Create a (name: default_value) map from flags.param_specs. default_kwargs = { name: flags.param_specs[name].default_value for name in flags.param_specs } params = Params(**default_kwargs)._replace(**kwargs) validate_params(params) return params def make_params_from_flags(): """Create a Params tuple for BenchmarkCNN from absl_flags.FLAGS. Returns: Params namedtuple for constructing BenchmarkCNN. """ # Collect (name: value) pairs for absl_flags.FLAGS with matching names in # flags.param_specs. flag_values = {name: getattr(absl_flags.FLAGS, name) for name in flags.param_specs.keys()} return Params(**flag_values) def remove_param_fields(params, fields_to_remove): """Remove fields from a Params namedtuple.""" params_dict = params._asdict() for field in fields_to_remove: assert field in params_dict, 'Invalid Params field: ' + field params_dict = {k: v for k, v in params_dict.items() if k not in fields_to_remove} new_params_type = namedtuple('Params', params_dict.keys()) return new_params_type(**params_dict) def get_num_batches_and_epochs(params, batch_size, num_examples_per_epoch): """Returns the number of batches and epochs to run for. Args: params: Params tuple, typically created by make_params or make_params_from_flags. batch_size: The number of images per step. num_examples_per_epoch: The number of images in a single epoch. Returns: num_batches: The number of batches to run for. num_epochs: The number of epochs to run for. This might be slightly smaller than params.num_epochs if specified, because the number of batches must be an integer. Raises: ValueError: Invalid or unsupported params. """ if params.num_batches and params.num_epochs: raise ValueError('At most one of --num_batches and --num_epochs may be ' 'specified.') if params.num_epochs: num_batches = int(params.num_epochs * num_examples_per_epoch + batch_size - 1) // batch_size else: num_batches = params.num_batches or _DEFAULT_NUM_BATCHES num_epochs = num_batches * batch_size / num_examples_per_epoch return (num_batches, num_epochs) def get_piecewise_learning_rate(piecewise_learning_rate_schedule, global_step, num_batches_per_epoch): """Returns a piecewise learning rate tensor. Args: piecewise_learning_rate_schedule: The --piecewise_learning_rate_schedule parameter global_step: Scalar tensor representing the global step. num_batches_per_epoch: float indicating the number of batches per epoch. Returns: A scalar float tensor, representing the learning rate. Raises: ValueError: piecewise_learning_rate_schedule is not formatted correctly. """ pieces = piecewise_learning_rate_schedule.split(';') if len(pieces) % 2 == 0: raise ValueError('--piecewise_learning_rate_schedule must have an odd ' 'number of components') values = [] boundaries = [] for i, piece in enumerate(pieces): if i % 2 == 0: try: values.append(float(piece)) except ValueError: raise ValueError('Invalid learning rate: ' + piece) else: try: boundaries.append(int(int(piece) * num_batches_per_epoch) - 1) except ValueError: raise ValueError('Invalid epoch: ' + piece) return tf.train.piecewise_constant(global_step, boundaries, values, name='piecewise_learning_rate') def get_learning_rate(params, global_step, num_examples_per_epoch, model, batch_size): """Returns a learning rate tensor based on global_step. Args: params: Params tuple, typically created by make_params or make_params_from_flags. global_step: Scalar tensor representing the global step. num_examples_per_epoch: The number of examples per epoch. model: The model.Model object to obtain the default learning rate from if no learning rate is specified. batch_size: Number of examples per step Returns: A scalar float tensor, representing the learning rate. When evaluated, the learning rate depends on the current value of global_step. Raises: ValueError: Invalid or unsupported params. """ with tf.name_scope('learning_rate'): num_batches_per_epoch = num_examples_per_epoch / batch_size if params.piecewise_learning_rate_schedule: if (params.init_learning_rate is not None or params.learning_rate_decay_factor or params.minimum_learning_rate or params.num_epochs_per_decay): raise ValueError('No other learning rate-related flags can be ' 'specified if --piecewise_learning_rate_schedule is ' 'specified') learning_rate = get_piecewise_learning_rate( params.piecewise_learning_rate_schedule, global_step, num_batches_per_epoch) elif params.init_learning_rate is not None: learning_rate = params.init_learning_rate if (params.num_epochs_per_decay > 0 and params.learning_rate_decay_factor > 0): decay_steps = int(num_batches_per_epoch * params.num_epochs_per_decay) # Decay the learning rate exponentially based on the number of steps. learning_rate = tf.train.exponential_decay( params.init_learning_rate, global_step, decay_steps, params.learning_rate_decay_factor, staircase=True) if params.minimum_learning_rate != 0.: learning_rate = tf.maximum(learning_rate, params.minimum_learning_rate) else: learning_rate = model.get_learning_rate(global_step, batch_size) if params.num_learning_rate_warmup_epochs > 0 and ( params.init_learning_rate is not None or params.piecewise_learning_rate_schedule): warmup_steps = int(num_batches_per_epoch * params.num_learning_rate_warmup_epochs) init_lr = params.init_learning_rate if init_lr is None: init_lr = float(params.piecewise_learning_rate_schedule.split(';')[0]) warmup_lr = init_lr * tf.cast(global_step, tf.float32) / tf.cast( warmup_steps, tf.float32) learning_rate = tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: learning_rate) learning_rate = mlperf.logger.log_deferred_tensor_value( mlperf.tags.OPT_LR, learning_rate, global_step, every_n=100) return learning_rate def get_optimizer(params, learning_rate): """Returns the optimizer that should be used based on params.""" if params.optimizer == 'momentum': mlperf.logger.log(key=mlperf.tags.OPT_NAME, value=mlperf.tags.SGD_WITH_MOMENTUM) mlperf.logger.log(key=mlperf.tags.OPT_MOMENTUM, value=params.momentum) opt = tf.train.MomentumOptimizer( learning_rate, params.momentum, use_nesterov=True) elif params.optimizer == 'sgd': mlperf.logger.log(key=mlperf.tags.OPT_NAME, value=mlperf.tags.SGD) opt = tf.train.GradientDescentOptimizer(learning_rate) elif params.optimizer == 'rmsprop': opt = tf.train.RMSPropOptimizer( learning_rate, params.rmsprop_decay, momentum=params.rmsprop_momentum, epsilon=params.rmsprop_epsilon) elif params.optimizer == 'adam': opt = tf.train.AdamOptimizer(learning_rate, params.adam_beta1, params.adam_beta2, params.adam_epsilon) else: raise ValueError('Optimizer "{}" was not recognized'. format(params.optimizer)) return opt def generate_tfprof_profile(profiler, tfprof_file): """Generates a tfprof profile, writing it to a file and printing top ops. Args: profiler: A tf.profiler.Profiler. `profiler.add_step` must have already been called. tfprof_file: The filename to write the ProfileProto to. """ profile_proto = profiler.serialize_to_string() log_fn('Dumping ProfileProto to %s' % tfprof_file) with gfile.Open(tfprof_file, 'wb') as f: f.write(profile_proto) # Print out the execution times of the top operations. Note this # information can also be obtained with the dumped ProfileProto, but # printing it means tfprof doesn't have to be used if all the user wants # is the top ops. options = tf.profiler.ProfileOptionBuilder.time_and_memory() options['max_depth'] = _NUM_OPS_TO_PRINT options['order_by'] = 'accelerator_micros' profiler.profile_operations(options) class BenchmarkCNN(object): """Class for benchmarking a cnn network.""" def __init__(self, params, dataset=None, model=None): """Initialize BenchmarkCNN. Args: params: Params tuple, typically created by make_params or make_params_from_flags. dataset: If not None, the dataset to use. Otherwise, params is used to obtain the dataset. model: If not None, the model to use. Otherwise, params is used to obtain the model. Raises: ValueError: Unsupported params settings. """ mlperf.logger.log(key=mlperf.tags.RUN_START) self.params = params if params.eval: self._doing_eval = True else: # Note self._doing_eval can later switch to True in self._do_eval() if # self.params.eval_during_training_* is specified. self._doing_eval = False self.dataset = dataset or datasets.create_dataset(self.params.data_dir, self.params.data_name) self.model = model or model_config.get_model_config( self.params.model, self.dataset, self.params) self.trace_filename = self.params.trace_file self.rewriter_config = self.params.rewriter_config autotune_threshold = self.params.autotune_threshold if ( self.params.autotune_threshold) else 1 min_autotune_warmup = 5 * autotune_threshold * autotune_threshold self.num_warmup_batches = self.params.num_warmup_batches if ( self.params.num_warmup_batches is not None) else max( 10, min_autotune_warmup) self.graph_file = self.params.graph_file self.resize_method = self.params.resize_method self.sync_queue_counter = 0 self.num_gpus = self.params.num_gpus if self.params.gpu_indices: self.gpu_indices = [int(x) for x in self.params.gpu_indices.split(',')] else: self.gpu_indices = [x for x in range(self.num_gpus)] if (self.params.device == 'cpu' and self.params.data_format == 'NCHW' and not self.params.mkl): raise ValueError('device=cpu requires that data_format=NHWC') if ((self.params.num_epochs_per_decay or self.params.learning_rate_decay_factor) and not (self.params.init_learning_rate is not None and self.params.num_epochs_per_decay and self.params.learning_rate_decay_factor)): raise ValueError('If one of num_epochs_per_decay or ' 'learning_rate_decay_factor is set, both must be set' 'and learning_rate must be set') if (self.params.minimum_learning_rate and not (self.params.init_learning_rate is not None and self.params.num_epochs_per_decay and self.params.learning_rate_decay_factor)): raise ValueError('minimum_learning_rate requires learning_rate,' 'num_epochs_per_decay, and ' 'learning_rate_decay_factor to be set') if (self.params.use_fp16 and self.params.fp16_vars and 'replicated' in self.params.variable_update and self.params.all_reduce_spec and 'nccl' in self.params.all_reduce_spec): raise ValueError('fp16 variables are not supported with NCCL') if (self.params.use_fp16 and self.params.fp16_vars and self.params.gradient_repacking): raise ValueError('--fp16_vars cannot be used with --gradient_repacking') if self.params.variable_update == 'horovod' and self.params.num_gpus > 1: raise ValueError('Horovod benchmarks require num_gpus=1 on each worker') if self.params.variable_update == 'horovod' and self.params.job_name: raise ValueError('job_name should not be specified for Horovod.') if self.params.use_fp16 and self.params.fp16_enable_auto_loss_scale: if self.params.all_reduce_spec and 'nccl' in self.params.all_reduce_spec: raise ValueError('Automatic loss scaling is not supported with NCCL.') if self.params.variable_update not in ('parameter_server', 'replicated', 'independent'): raise ValueError('Automatic loss scaling is not supported with ' 'variable_update=%s.' % self.params.variable_update) if self.params.staged_vars: raise ValueError('Automatic loss scaling is not supported with' 'staged_vars.') if (self.params.debugger is not None and self.params.debugger != 'cli' and ':' not in self.params.debugger): raise ValueError('--debugger must be "cli" or in the form ' 'host:port') if self.params.hierarchical_copy and self.params.num_gpus <= 1: raise ValueError('--hierarchical_copy requires --num_gpus to be greater ' 'than 1') if params.save_model_secs and params.save_model_steps: raise ValueError('At most one of --save_model_secs and ' '--save_model_steps can be specified') eval_during_training_flags = list(map(bool, [ params.eval_during_training_every_n_steps, params.eval_during_training_every_n_epochs, params.eval_during_training_at_specified_steps, params.eval_during_training_at_specified_epochs, ])) if eval_during_training_flags.count(True) > 1: raise ValueError('At most one flag with --eval_during_training_* prefix ' 'must be specified.') eval_during_training_enabled = any(eval_during_training_flags) if eval_during_training_enabled: if params.eval: raise ValueError('At most one of --eval and --eval_during_training_* ' 'must be specified') if params.forward_only: raise ValueError('At most one of --forward_only and ' '--eval_during_training_* must be specified') if params.job_name: raise ValueError('--eval_during_training_* is not yet supported in ' 'distributed mode.') if params.staged_vars: raise ValueError('--eval_during_training_* is not currently compatible ' 'with --staged_vars') if params.stop_at_top_1_accuracy and not eval_during_training_enabled: raise ValueError('--stop_at_top_1_accuracy is only supported with ' '--eval_during_training_*') if params.collect_eval_results_async and params.model != 'ssd300': raise ValueError('--collect_eval_results_async only works with ssd300 ' 'model currently.') if self.params.forward_only and self.params.freeze_when_forward_only: if self.params.train_dir is not None: raise ValueError('In forward_only mode, when --freeze_when_forward_only' ' is True, --train_dir should not be specified') if self.params.data_dir and not self.params.datasets_use_prefetch: raise ValueError('In forward_only mode, when --freeze_when_forward_only' ' is True and --data_dir is set, ' '--datasets_use_prefetch should be set to True') if self.params.job_name: raise ValueError('In forward_only mode, when --freeze_when_forward_only' ' is True, --job_name should not be specified and ' 'distributed running is not supported') self.forward_only_and_freeze = True else: self.forward_only_and_freeze = False if self.params.trt_mode: raise ValueError('--trt_mode should not be specified if one of ' '--forward_only and --freeze_when_forward_only is set ' 'to False') self.mode = get_mode_from_params(self.params) # Use the batch size from the command line if specified, otherwise use the # model's default batch size. Scale the benchmark's batch size by the # number of GPUs. if self.params.batch_size > 0: self.model.set_batch_size(self.params.batch_size) self.batch_size = self.model.get_batch_size() * self.num_gpus if self.mode in (constants.BenchmarkMode.TRAIN, constants.BenchmarkMode.TRAIN_AND_EVAL): self.train_batch_size = self.batch_size else: self.train_batch_size = None if self.mode in (constants.BenchmarkMode.EVAL, constants.BenchmarkMode.TRAIN_AND_EVAL): if self.params.eval_batch_size > 0: self.eval_batch_size = self.params.eval_batch_size * self.num_gpus else: self.eval_batch_size = self.batch_size else: self.eval_batch_size = None self.batch_group_size = self.params.batch_group_size self.enable_auto_loss_scale = ( self.params.use_fp16 and self.params.fp16_enable_auto_loss_scale) self.loss_scale = None self.loss_scale_normal_steps = None self.job_name = self.params.job_name # "" for local training # PS server is used for distributed jobs not using all-reduce. use_ps_server = self.job_name and (self.params.variable_update != 'distributed_all_reduce' and self.params.variable_update != 'collective_all_reduce') # controller is used for distributed_all_reduce with > 1 worker. use_controller = ( self.params.variable_update == 'distributed_all_reduce' and self.job_name) if use_controller and not params.controller_host: raise ValueError('When variable_update==distributed_all_reduce ' 'controller_host must also be specified.') self.single_session = ( self.params.variable_update == 'distributed_all_reduce') # collective_all_reduce doesn't need a controller or ps self.distributed_collective = ( self.params.variable_update == 'collective_all_reduce' and self.job_name) self.local_parameter_device_flag = self.params.local_parameter_device if self.job_name: self.task_index = self.params.task_index self.cluster_manager = platforms_util.get_cluster_manager( params, create_config_proto(params)) assert isinstance(self.cluster_manager, cnn_util.BaseClusterManager) worker_prefix = '/job:worker/replica:0/task:%s' % self.task_index if use_ps_server: self.param_server_device = tf.train.replica_device_setter( worker_device=worker_prefix + '/cpu:0', cluster=self.cluster_manager.get_cluster_spec()) # This device on which the queues for managing synchronization between # servers should be stored. self.sync_queue_devices = [ '/job:ps/replica:0/task:%s/cpu:0' % i for i in range(self.cluster_manager.num_ps()) ] else: self.sync_queue_devices = ['/job:worker/replica:0/task:0/cpu:0'] else: self.task_index = 0 self.cluster_manager = None worker_prefix = '' self.param_server_device = '/%s:0' % self.params.local_parameter_device self.sync_queue_devices = [self.param_server_device] if self.cluster_manager: self.num_workers = self.cluster_manager.num_workers() elif self.params.variable_update == 'horovod': import horovod.tensorflow as hvd # pylint: disable=g-import-not-at-top self.num_workers = hvd.size() else: self.num_workers = 1 self.num_ps = self.cluster_manager.num_ps() if self.cluster_manager else 0 if self.num_workers > 1 and self.params.all_reduce_spec == 'nccl': raise ValueError('--all_reduce_spec=nccl is invalid in a ' 'multi-worker job') # Device to use for ops that need to always run on the local worker's CPU. self.cpu_device = '%s/cpu:0' % worker_prefix # Device to use for ops that need to always run on the local worker's # 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) ] subset = 'validation' if params.eval else 'train' self.num_batches, self.num_epochs = get_num_batches_and_epochs( params, self.batch_size * self.num_workers, self.dataset.num_examples_per_epoch(subset)) if self.mode in (constants.BenchmarkMode.EVAL, constants.BenchmarkMode.TRAIN_AND_EVAL): # TODO(reedwm): Currently we do extra eval logic for num_eval_batches and # the preprocessor. We should encapsulate this logic into a shared # function or class. if params.num_eval_batches is None and params.num_eval_epochs is None: eval_params = self.params else: eval_params = self.params._replace( num_batches=self.params.num_eval_batches, num_epochs=self.params.num_eval_epochs) self.num_eval_batches, self.num_eval_epochs = get_num_batches_and_epochs( eval_params, self.eval_batch_size * self.num_workers, self.dataset.num_examples_per_epoch('validation')) else: self.num_eval_batches, self.num_eval_epochs = None, None num_train_examples_per_epoch = self.dataset.num_examples_per_epoch('train') if self.params.eval_during_training_every_n_epochs: n_epochs = self.params.eval_during_training_every_n_epochs self.eval_during_training_at_specified_steps = { (int(e * num_train_examples_per_epoch + self.batch_size - 1) // self.batch_size) for e in np.arange(n_epochs, self.num_epochs, n_epochs)} if self.params.eval_during_training_at_specified_steps: try: self.eval_during_training_at_specified_steps = set(map( int, self.params.eval_during_training_at_specified_steps)) except ValueError: raise ValueError('Param eval_during_training_at_specified_steps value ' 'of %s cannot be converted to a list of integers.' % (self.params.eval_during_training_at_specified_steps)) if self.params.eval_during_training_at_specified_epochs: try: n_epochs = list(map( float, self.params.eval_during_training_at_specified_epochs)) offset = n_epochs[0] - 1 if offset.is_integer(): offset = int(offset) mlperf.logger.log(key=mlperf.tags.EVAL_EPOCH_OFFSET, value=offset) self.eval_during_training_at_specified_steps = { (int(e * num_train_examples_per_epoch + self.batch_size - 1) // self.batch_size) for e in n_epochs} except ValueError: raise ValueError('Param eval_during_training_at_specified_epochs value ' 'of %s cannot be converted to a list of floats.' % (self.params.eval_during_training_at_specified_epochs)) if params.eval_during_training_every_n_epochs: offset = params.eval_during_training_every_n_epochs - 1 if offset.is_integer(): offset = int(offset) mlperf.logger.log(key=mlperf.tags.EVAL_EPOCH_OFFSET, value=offset) if (self.params.staged_vars and self.params.variable_update != 'parameter_server'): raise ValueError('staged_vars for now is only supported with ' 'variable_update=parameter_server') if self.params.variable_update == 'parameter_server': if self.job_name: if not self.params.staged_vars: self.variable_mgr = variable_mgr.VariableMgrDistributedFetchFromPS( self) else: self.variable_mgr = ( variable_mgr.VariableMgrDistributedFetchFromStagedPS(self)) else: if not self.params.staged_vars: self.variable_mgr = variable_mgr.VariableMgrLocalFetchFromPS(self) else: self.variable_mgr = variable_mgr.VariableMgrLocalFetchFromStagedPS( self) elif self.params.variable_update == 'replicated': if self.job_name: raise ValueError('Invalid variable_update in distributed mode: %s' % self.params.variable_update) self.variable_mgr = variable_mgr.VariableMgrLocalReplicated( self, self.params.all_reduce_spec, self.params.agg_small_grads_max_bytes, self.params.agg_small_grads_max_group, self.params.allreduce_merge_scope) elif self.params.variable_update == 'distributed_all_reduce': assert self.params.cross_replica_sync self.variable_mgr = variable_mgr.VariableMgrDistributedAllReduce( self, self.params.all_reduce_spec, ('worker' if self.num_workers > 1 else 'localhost'), self.num_workers, self.params.agg_small_grads_max_bytes, self.params.agg_small_grads_max_group, self.params.allreduce_merge_scope) elif self.params.variable_update == 'collective_all_reduce': assert self.params.cross_replica_sync self.variable_mgr = variable_mgr.VariableMgrCollectiveAllReduce( self, self.params.all_reduce_spec, self.num_workers, self.num_gpus, self.task_index, self.params.allreduce_merge_scope) elif self.params.variable_update == 'distributed_replicated': assert self.params.cross_replica_sync if not self.job_name: raise ValueError('Invalid variable_update in local mode: %s' % self.params.variable_update) self.variable_mgr = variable_mgr.VariableMgrDistributedReplicated(self) elif self.params.variable_update in ('independent', 'horovod'): if self.job_name: raise ValueError('Invalid variable_update in distributed mode: %s' % self.params.variable_update) self.variable_mgr = variable_mgr.VariableMgrIndependent(self) else: raise ValueError( 'Invalid variable_update: %s' % self.params.variable_update) # Device to use for running on the local worker's compute device, but # with variables assigned to parameter server devices. self.devices = self.variable_mgr.get_devices() if self.job_name: if use_ps_server: self.global_step_device = self.param_server_device elif self.params.variable_update == 'collective_all_reduce': self.global_step_device = self.cpu_device else: self.global_step_device = '/job:worker/replica:0/task:0/cpu:0' else: self.global_step_device = self.cpu_device self.input_preprocessor = None self.eval_input_preprocessor = None if not self.dataset.use_synthetic_gpu_inputs(): if not self.params.eval: self.input_preprocessor = self.get_input_preprocessor() if self.mode in (constants.BenchmarkMode.EVAL, constants.BenchmarkMode.TRAIN_AND_EVAL): with self._do_eval(): self.eval_input_preprocessor = self.get_input_preprocessor() self.datasets_use_prefetch = ( self.params.datasets_use_prefetch and # TODO(rohanj): Figure out why --datasets_use_prefetch freezes on the # CPU. self.params.device.lower() != 'cpu' and self.input_preprocessor and self.input_preprocessor.supports_datasets()) self.init_global_step = 0 self._config_benchmark_logger() if self.mode == constants.BenchmarkMode.TRAIN_AND_EVAL: # Remove "eval" from params so it is not accidentally used. Since eval can # still occur despite params.eval being False, params.eval should never # be used. We cannot yet remove this unconditionally, because the SSD # model still uses params.eval, and hence does not work properly with # --eval_during_training_*. # TODO(b/116627045): We should also remove fields that have an eval # equivalent, like num_batches and num_eval_batches. self.params = remove_param_fields(self.params, {'eval'}) @contextlib.contextmanager def _do_eval(self): """Context manager to switches BenchmarkCNN to eval mode. Any evaluation code should be put under this context manager. This context manager switches self._doing_eval to True. It also switches certain attributes, like self.num_batches and self.num_epochs, to be the number of batches and epochs for evaluation respectively Yields: Nothing. """ # TODO(b/116627045): Find a more general way of switching attributes to the # eval equivalents. old_doing_eval = self._doing_eval old_num_batches = self.num_batches old_num_epochs = self.num_epochs old_batch_size = self.batch_size try: self._doing_eval = True self.num_batches = self.num_eval_batches self.num_epochs = self.num_eval_epochs self.batch_size = self.eval_batch_size self.model.set_batch_size(self.eval_batch_size // self.num_gpus) yield finally: self._doing_eval = old_doing_eval self.num_batches = old_num_batches self.num_epochs = old_num_epochs self.batch_size = old_batch_size self.model.set_batch_size(old_batch_size // self.num_gpus) def _config_benchmark_logger(self): """Config the model garden benchmark logger.""" model_benchmark_logger = None if self.params.benchmark_log_dir is not None: try: from official.r1.utils.logs import logger as models_logger # pylint: disable=g-import-not-at-top except ImportError: tf.logging.fatal('Please include tensorflow/models to the PYTHONPATH ' 'in order to use BenchmarkLogger. Configured ' 'benchmark_log_dir: %s' % self.params.benchmark_log_dir) raise model_benchmark_logger = models_logger.BenchmarkFileLogger( self.params.benchmark_log_dir) self.benchmark_logger = model_benchmark_logger # TODO(laigd): this changes the global device list which is used everywhere, # consider refactoring it. def reset_devices_for_task(self, task_num, is_local=False): """Used to imitate another task when building a distributed graph.""" worker_prefix = ('/job:localhost' if is_local else '/job:worker/replica:0/task:%s' % task_num) 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) ] self.devices = self.variable_mgr.get_devices() def raw_devices_across_tasks(self, is_local=False): """Returns list of raw device names across all tasks.""" if is_local: assert self.num_workers == 1 return self.raw_devices 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) ] def print_info(self): """Print basic information.""" benchmark_info = self._get_params_info() log_fn('Model: %s' % self.model.get_model_name()) log_fn('Dataset: %s' % benchmark_info['dataset_name']) log_fn('Mode: %s' % self.mode) log_fn('SingleSess: %s' % benchmark_info['single_session']) log_fn('Batch size: %s global' % (self.batch_size * self.num_workers)) log_fn(' %s per device' % (self.batch_size // len(self.raw_devices))) if self.batch_group_size > 1: log_fn(' %d batches per prepocessing group' % self.batch_group_size) log_fn('Num batches: %d' % self.num_batches) log_fn('Num epochs: %.2f' % self.num_epochs) log_fn('Devices: %s' % benchmark_info['device_list']) log_fn('NUMA bind: %s' % self.params.use_numa_affinity) log_fn('Data format: %s' % self.params.data_format) if self.rewriter_config: log_fn('RewriterConfig: %s' % self.rewriter_config) log_fn('Optimizer: %s' % self.params.optimizer) log_fn('Variables: %s' % self.params.variable_update) if (self.params.variable_update == 'replicated' or self.params.variable_update == 'distributed_all_reduce' or self.params.variable_update == 'collective_all_reduce'): log_fn('AllReduce: %s' % self.params.all_reduce_spec) if self.job_name: log_fn('Sync: %s' % self.params.cross_replica_sync) if self.params.staged_vars: log_fn('Staged vars: %s' % self.params.staged_vars) if self.params.variable_update == 'horovod' and self.params.horovod_device: log_fn('Horovod on: %s' % self.params.horovod_device) log_fn('==========') def _get_params_info(self): """Get the common parameters info for the benchmark run. Returns: A dict of processed parameters. """ dataset_name = self.dataset.name if self.dataset.use_synthetic_gpu_inputs(): dataset_name += ' (synthetic)' single_session = self.params.variable_update == 'distributed_all_reduce' if single_session: device_list = self.raw_devices_across_tasks() elif self.params.variable_update == 'horovod': device_list = ['horovod/%s:%d' % (self.params.device, idx) for idx in range(self.num_workers)] else: device_list = self.raw_devices return { 'dataset_name': dataset_name, 'single_session': single_session, 'device_list': device_list,} def _log_benchmark_run(self): """Log the benchmark info to the logger. The info logged here should be similar to print_info(), but in a structured JSON format. """ if self.benchmark_logger: benchmark_info = self._get_params_info() run_param = { 'model': self.model.get_model_name(), 'dataset': benchmark_info['dataset_name'], 'mode': self.mode, 'single_sess': benchmark_info['single_session'], 'devices': benchmark_info['device_list'], 'batch_size': self.batch_size, 'batch_size_per_device': self.batch_size // len(self.raw_devices), 'num_batches':
{'commit': 'diff --git a/scripts/tf_cnn_benchmarks/benchmark_cnn.py b/scripts/tf_cnn_benchmarks/benchmark_cnn.py\nindex ba64dd6..f5b837b 100644\n--- a/scripts/tf_cnn_benchmarks/benchmark_cnn.py\n+++ b/scripts/tf_cnn_benchmarks/benchmark_cnn.py\n@@ -2124,7 +2124,7 @@ class BenchmarkCNN(object):\n         images = tf.reshape(\n             images,\n             shape=[\n-                self.batch_size, image_size, image_size,\n+                self.batch_size // self.num_gpus, image_size, image_size,\n                 self.dataset.depth\n             ])\n     else:\n', 'commit_subject': 'Fixing the per-gpu batch size.', 'repo_name': 'tensorflow/benchmarks', 'committer': {'name': 'Reed', 'email': '75734eae9db3e4fb4d6c671d22722b468d1a389b@google.com', 'time_sec': '1521849283', 'tz_offset': '-420', 'date': {'seconds': '1521849283'}}, 'commit_tree': 'd06f50f88b74028be893ca41fcc45e873b529362', 'commit_message': 'Fixing the per-gpu batch size.\n\nPiperOrigin-RevId: 190181787\n', 'language_name': 'Shell', 'license': 'apache-2.0', 'hunks': '@@ -2124,7 +2124,7 @@ class BenchmarkCNN(object):\n         images = tf.reshape(\n             images,\n             shape=[\n-                self.batch_size, image_size, image_size,\n+                self.batch_size // self.num_gpus, image_size, image_size,\n                 self.dataset.depth\n             ])\n     else:\n', 'addition_count': 1, 'deletion_count': 1, 'src_file': 'a/scripts/tf_cnn_benchmarks/benchmark_cnn.py', 'tgt_file': 'b/scripts/tf_cnn_benchmarks/benchmark_cnn.py', 'file_extension': '.py', 'before_file': '# Copyright 2017 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the "License");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an "AS IS" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n"""TensorFlow benchmark library.\n\nSee the README for more information.\n"""\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport argparse\nfrom collections import namedtuple\nimport contextlib\nimport math\nimport multiprocessing\nimport os\nimport re\nimport threading\nimport time\nimport traceback\n\nfrom absl import flags as absl_flags\nimport numpy as np\n\nimport six\nfrom six.moves import xrange  # pylint: disable=redefined-builtin\nimport tensorflow.compat.v1 as tf\n\n# pylint: disable=g-direct-tensorflow-import\nimport cnn_util\nimport constants\nimport datasets\nimport flags\nimport mlperf\nimport variable_mgr\nimport variable_mgr_util\nfrom cnn_util import log_fn\nfrom models import model_config\nfrom platforms import util as platforms_util\nfrom google.protobuf import text_format\nfrom tensorflow.core.protobuf import rewriter_config_pb2\nfrom tensorflow.python import debug as tf_debug\nfrom tensorflow.python.client import timeline\nfrom tensorflow.python.framework import graph_util\nfrom tensorflow.python.framework import graph_util_impl\nfrom tensorflow.python.framework import importer\nfrom tensorflow.python.ops import data_flow_ops\nfrom tensorflow.python.platform import gfile\nfrom tensorflow.python.util import nest\n\n\n_DEFAULT_NUM_BATCHES = 100\n\n\n# GraphInfo encapsulates the tensors/ops that we care about after building a\n# graph. We use them to benchmark the graph.\nGraphInfo = namedtuple(  # pylint: disable=invalid-name\n    \'GraphInfo\',\n    [\n        # Ops that produce the input batches (before preprocessing).\n        \'input_producer_op\',\n        # Ops that adds the preprocessed images to the staging areas\n        \'enqueue_ops\',\n        # Fetches of sess.run()\n        \'fetches\',\n        # Op that performs synchronization in distributed mode\n        \'execution_barrier\',\n        # The global step variable\n        \'global_step\',\n        # Group of ops that perform per-device initialization work\n        \'local_var_init_op_group\',\n        # Op to produce summaries\n        \'summary_op\'\n    ])\n\n\n# InputProcessingInfo contains various sources of inputs which will be later fed\n# into the model. If synthetic data is used, all three fields are None.\nInputProcessingInfo = namedtuple(\n    \'InputProcessingInfo\',\n    [\n        # The first two fields are non-None iff datasets prefetching is not\n        # used.\n\n        # Ops that produce the input batches.\n        \'input_producer_op\',\n        # A list of StagingArea for each device.\n        \'input_producer_stages\',\n\n        # Input produced using multi device iterator. Non-None iff datasets\n        # prefetching is used\n        \'multi_device_iterator_input\'\n    ])\n\n\n# TODO(reedwm): add upper_bound and lower_bound to appropriate integer and\n# float flags, and change certain string flags to enum flags.\n\nflags.DEFINE_string(\'model\', \'trivial\',\n                    \'Name of the model to run, the list of supported models \'\n                    \'are defined in models/model.py\')\n# The code will first check if it\'s running under benchmarking mode\n# or evaluation mode, depending on \'eval\':\n# Under the evaluation mode, this script will read a saved model,\n#   and compute the accuracy of the model against a validation dataset.\n#   Additional ops for accuracy and top_k predictors are only used under\n#   this mode.\n# Under the benchmarking mode, user can specify whether nor not to use\n#   the forward-only option, which will only compute the loss function.\n#   forward-only cannot be enabled with eval at the same time.\nflags.DEFINE_boolean(\'eval\', False, \'whether use eval or benchmarking\')\nflags.DEFINE_integer(\'eval_interval_secs\', 0,\n                     \'How often to run eval on saved checkpoints. Usually the \'\n                     \'same as save_model_secs from the corresponding training \'\n                     \'run. Pass 0 to eval only once.\')\nflags.DEFINE_integer(\'eval_during_training_every_n_steps\', None,\n                     \'Every n steps during training, pause training, run \'\n                     \'evaluation, then resume training. Must not be used with \'\n                     \'--eval, as unlike --eval, this option causes both \'\n                     \'training and eval to be done. This may take slightly \'\n                     \'more GPU memory than running just training or evaluation \'\n                     \'alone. It also may slightly slow down training, even \'\n                     \'when not taking into account the additional time to \'\n                     \'evaluate.\', lower_bound=1)\nflags.DEFINE_float(\'eval_during_training_every_n_epochs\', None,\n                   \'After every n training epochs, pause training, run \'\n                   \'evaluation, then resume training. See \'\n                   \'--eval_during_training_every_n_steps for more information.\')\nflags.DEFINE_list(\'eval_during_training_at_specified_steps\', [],\n                  \'Specify a list of training steps, pause training at each of \'\n                  \'these steps, run evaluation, then resume training. See \'\n                  \'--eval_during_training_every_n_steps for more information.\')\nflags.DEFINE_list(\'eval_during_training_at_specified_epochs\', [],\n                  \'Specify a list of training epochs, pause training after \'\n                  \'each of these epochs, run evaluation, then resume training. \'\n                  \'See --eval_during_training_every_n_steps for more \'\n                  \'information.\')\nflags.DEFINE_boolean(\'forward_only\', False,\n                     \'whether use forward-only or training for benchmarking\')\nflags.DEFINE_boolean(\'freeze_when_forward_only\', False,\n                     \'whether to freeze the graph when in forward-only mode.\')\nflags.DEFINE_boolean(\'print_training_accuracy\', False,\n                     \'whether to calculate and print training accuracy during \'\n                     \'training\')\nflags.DEFINE_integer(\'batch_size\', 0, \'batch size per compute device\')\nflags.DEFINE_integer(\'eval_batch_size\', 0, \'eval batch size per compute device\')\nflags.DEFINE_integer(\'batch_group_size\', 1,\n                     \'number of groups of batches processed in the image \'\n                     \'producer.\')\nflags.DEFINE_integer(\'num_batches\', None, \'number of batches to run, excluding \'\n                     \'warmup. Defaults to %d\' % _DEFAULT_NUM_BATCHES)\nflags.DEFINE_integer(\'num_eval_batches\', None,\n                     \'number of eval batches to run, excluding warmup. \'\n                     \'Defaults to --num_batches\')\nflags.DEFINE_float(\'num_epochs\', None,\n                   \'number of epochs to run, excluding warmup. \'\n                   \'This and --num_batches cannot both be specified.\')\nflags.DEFINE_float(\'num_eval_epochs\', None,\n                   \'number of eval epochs to run, excluding warmup. \'\n                   \'Defaults to --num_epochs\')\nflags.DEFINE_float(\'stop_at_top_1_accuracy\', None,\n                   \'If set, stops training after the evaluation accuracy hits \'\n                   \'this number. Can only be used with one of the \'\n                   \'--eval_during_training_* flags.\')\nflags.DEFINE_boolean(\'collect_eval_results_async\', False,\n                     \'If True, start a separate process to postprocess eval \'\n                     \'results asynchronously. This currently only works with \'\n                     \'the SSD model.\')\nflags.DEFINE_integer(\'num_warmup_batches\', None,\n                     \'number of batches to run before timing\')\nflags.DEFINE_integer(\'autotune_threshold\', None,\n                     \'The autotune threshold for the models\')\n# TODO(tucker): change num_gpus to num_devices\nflags.DEFINE_integer(\'num_gpus\', 1, \'the number of GPUs to run on\')\nflags.DEFINE_string(\'gpu_indices\', \'\', \'indices of worker GPUs in ring order\')\nflags.DEFINE_integer(\'display_every\', 10,\n                     \'Number of local steps after which progress is printed \'\n                     \'out\')\nflags.DEFINE_float(\'display_perf_ewma\', None,\n                   \'If set, display numbers of images/sec using exponentially \'\n                   \'weighted moving avearge with the specified weight, which \'\n                   \'defines how much current value contributes to the reported \'\n                   \'average. Increasing weight makes the reported performance \'\n                   \'number reflect more about the real-time speed instead of \'\n                   \'the entire history\', lower_bound=0, upper_bound=1)\nflags.DEFINE_string(\'data_dir\', None,\n                    \'Path to dataset in TFRecord format (aka Example \'\n                    \'protobufs). If not specified, synthetic data will be \'\n                    \'used.\')\nflags.DEFINE_string(\'data_name\', None,\n                    \'Name of dataset: imagenet or cifar10. If not specified, \'\n                    \'it is automatically guessed based on data_dir.\')\nflags.DEFINE_string(\'resize_method\', \'bilinear\',\n                    \'Method for resizing input images: crop, nearest, \'\n                    \'bilinear, bicubic, area, or round_robin. The `crop` mode \'\n                    \'requires source images to be at least as large as the \'\n                    \'network input size. The `round_robin` mode applies \'\n                    \'different resize methods based on position in a batch in \'\n                    \'a round-robin fashion. Other modes support any sizes and \'\n                    \'apply random bbox distortions before resizing (even with \'\n                    \'distortions=False).\')\nflags.DEFINE_boolean(\'distortions\', False,\n                     \'Enable/disable distortions during image preprocessing. \'\n                     \'These include bbox and color distortions.\')\nflags.DEFINE_boolean(\'use_datasets\', True,\n                     \'Enable use of datasets for input pipeline\')\nflags.DEFINE_string(\'input_preprocessor\', \'default\',\n                    \'Name of input preprocessor. The list of supported input \'\n                    \'preprocessors are defined in preprocessing.py.\')\nflags.DEFINE_string(\'gpu_thread_mode\', \'gpu_private\',\n                    \'Methods to assign GPU host work to threads. \'\n                    \'global: all GPUs and CPUs share the same global threads; \'\n                    \'gpu_private: a private threadpool for each GPU; \'\n                    \'gpu_shared: all GPUs share the same threadpool.\')\nflags.DEFINE_integer(\'per_gpu_thread_count\', 0,\n                     \'The number of threads to use for GPU. Only valid when \'\n                     \'gpu_thread_mode is not global.\')\nflags.DEFINE_boolean(\'hierarchical_copy\', False,\n                     \'Use hierarchical copies. Currently only optimized for \'\n                     \'use on a DGX-1 with 8 GPUs and may perform poorly on \'\n                     \'other hardware. Requires --num_gpus > 1, and only \'\n                     \'recommended when --num_gpus=8\')\n# TODO(hinsu): Support auto-detection of the network topology while still\n# retaining the ability to specify a particular topology for debugging.\nflags.DEFINE_enum(\n    \'network_topology\', constants.NetworkTopology.DGX1,\n    (constants.NetworkTopology.DGX1, constants.NetworkTopology.GCP_V100),\n    \'Network topology specifies the topology used to connect multiple devices. \'\n    \'Network topology is used to decide the hierarchy to use for the \'\n    \'hierarchical_copy.\')\nflags.DEFINE_integer(\'gradient_repacking\', 0, \'Use gradient repacking. It \'\n                     \'currently only works with replicated mode. At the end of \'\n                     \'each step, it repacks the gradients for more efficient \'\n                     \'cross-device transportation. A non-zero value specifies \'\n                     \'the number of split packs that will be formed.\',\n                     lower_bound=0)\nflags.DEFINE_boolean(\'compact_gradient_transfer\', True, \'Compact gradient\'\n                     \'as much as possible for cross-device transfer and \'\n                     \'aggregation.\')\nflags.DEFINE_enum(\'variable_consistency\', \'strong\', (\'strong\', \'relaxed\'),\n                  \'The data consistency for trainable variables. With strong \'\n                  \'consistency, the variable always have the updates from \'\n                  \'previous step. With relaxed consistency, all the updates \'\n                  \'will eventually show up in the variables. Likely one step \'\n                  \'behind.\')\nflags.DEFINE_boolean(\'datasets_repeat_cached_sample\', False,\n                     \'Enable use of a special datasets pipeline that reads a \'\n                     \'single TFRecord into memory and repeats it infinitely \'\n                     \'many times. The purpose of this flag is to make it \'\n                     \'possible to write regression tests that are not \'\n                     \'bottlenecked by CNS throughput. \'\n                     \'Use datasets_use_caching to cache input data.\')\nflags.DEFINE_enum(\'local_parameter_device\', \'gpu\', (\'cpu\', \'gpu\', \'CPU\', \'GPU\'),\n                  \'Device to use as parameter server: cpu or gpu. For \'\n                  \'distributed training, it can affect where caching of \'\n                  \'variables happens.\')\nflags.DEFINE_enum(\'device\', \'gpu\', (\'cpu\', \'gpu\', \'CPU\', \'GPU\'),\n                  \'Device to use for computation: cpu or gpu\')\nflags.DEFINE_enum(\'data_format\', \'NCHW\', (\'NHWC\', \'NCHW\'),\n                  \'Data layout to use: NHWC (TF native) or NCHW (cuDNN \'\n                  \'native, requires GPU).\')\nflags.DEFINE_integer(\'num_intra_threads\', None,\n                     \'Number of threads to use for intra-op parallelism. If \'\n                     \'set to 0, the system will pick an appropriate number. \'\n                     \'None is the same as 0 except that it disables intra-op \'\n                     \'parallelism on a GPU.\')\nflags.DEFINE_integer(\'num_inter_threads\', 0,\n                     \'Number of threads to use for inter-op parallelism. If \'\n                     \'set to 0, the system will pick an appropriate number.\')\nflags.DEFINE_boolean(\'use_numa_affinity\', False,\n                     \'Whether to turn on NUMA affinity for CPU devices. \'\n                     \'This is probably only useful when --device=cpu.\')\nflags.DEFINE_string(\'trace_file\', \'\',\n                    \'Enable TensorFlow tracing and write trace to this file.\')\nflags.DEFINE_boolean(\'use_chrome_trace_format\', True,\n                     \'If True, the trace_file, if specified, will be in a \'\n                     \'Chrome trace format. If False, then it will be a \'\n                     \'StepStats raw proto.\')\n_NUM_STEPS_TO_PROFILE = 10\n_NUM_OPS_TO_PRINT = 20\nflags.DEFINE_string(\'tfprof_file\', None,\n                    \'If specified, write a tfprof ProfileProto to this file. \'\n                    \'The performance and other aspects of the model can then \'\n                    \'be analyzed with tfprof. See \'\n                    \'https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/profiler/g3doc/command_line.md \'  # pylint: disable=line-too-long\n                    \'for more info on how to do this. The first %d steps \'\n                    \'are profiled. Additionally, the top %d most time \'\n                    \'consuming ops will be printed.\\n\'\n                    \'Note: profiling with tfprof is very slow, but most of the \'\n                    \'overhead is spent between steps. So, profiling results \'\n                    \'are more accurate than the slowdown would suggest.\' %\n                    (_NUM_STEPS_TO_PROFILE, _NUM_OPS_TO_PRINT))\nflags.DEFINE_string(\'graph_file\', None,\n                    \'Write the model\\\'s graph definition to this file. \'\n                    \'Defaults to binary format unless filename ends in "txt".\')\nflags.DEFINE_string(\'partitioned_graph_file_prefix\', None,\n                    \'If specified, after the graph has been partitioned and \'\n                    \'optimized, write out each partitioned graph to a file \'\n                    \'with the given prefix.\')\nflags.DEFINE_enum(\'optimizer\', \'sgd\', (\'momentum\', \'sgd\', \'rmsprop\', \'adam\'),\n                  \'Optimizer to use\')\nflags.DEFINE_float(\'init_learning_rate\', None,\n                   \'Initial learning rate for training.\')\nflags.DEFINE_string(\'piecewise_learning_rate_schedule\', None,\n                    \'Specifies a piecewise learning rate schedule based on the \'\n                    \'number of epochs. This is the form LR0;E1;LR1;...;En;LRn, \'\n                    \'where each LRi is a learning rate and each Ei is an epoch \'\n                    \'indexed from 0. The learning rate is LRi if the \'\n                    \'E(i-1) <= current_epoch < Ei. For example, if this \'\n                    \'paramater is 0.3;10;0.2;25;0.1, the learning rate is 0.3 \'\n                    \'for the first 10 epochs, then is 0.2 for the next 15 \'\n                    \'epochs, then is 0.1 until training ends.\')\nflags.DEFINE_float(\'num_epochs_per_decay\', 0,\n                   \'Steps after which learning rate decays. If 0, the learning \'\n                   \'rate does not decay.\')\nflags.DEFINE_float(\'learning_rate_decay_factor\', 0,\n                   \'Learning rate decay factor. Decay by this factor every \'\n                   \'`num_epochs_per_decay` epochs. If 0, learning rate does \'\n                   \'not decay.\')\nflags.DEFINE_float(\'num_learning_rate_warmup_epochs\', 0,\n                   \'Slowly increase to the initial learning rate in the first \'\n                   \'num_learning_rate_warmup_epochs linearly.\')\nflags.DEFINE_float(\'minimum_learning_rate\', 0,\n                   \'The minimum learning rate. The learning rate will \'\n                   \'never decay past this value. Requires `learning_rate`, \'\n                   \'`num_epochs_per_decay` and `learning_rate_decay_factor` to \'\n                   \'be set.\')\nflags.DEFINE_float(\'resnet_base_lr\', None, "Base learning rate at bs=256. Only "\n                   "relevant when training ResNet and utilizing the model\'s "\n                   "learning rate heuristic (get_learning_rate).")\nflags.DEFINE_float(\'momentum\', 0.9, \'Momentum for training.\')\nflags.DEFINE_float(\'rmsprop_decay\', 0.9, \'Decay term for RMSProp.\')\nflags.DEFINE_float(\'rmsprop_momentum\', 0.9, \'Momentum in RMSProp.\')\nflags.DEFINE_float(\'rmsprop_epsilon\', 1.0, \'Epsilon term for RMSProp.\')\nflags.DEFINE_float(\'adam_beta1\', 0.9, \'Beta2 term for the Adam optimizer\')\nflags.DEFINE_float(\'adam_beta2\', 0.999, \'Beta2 term for the Adam optimizer\')\nflags.DEFINE_float(\'adam_epsilon\', 1e-8, \'Epsilon term for the Adam optimizer\')\nflags.DEFINE_float(\'gradient_clip\', None,\n                   \'Gradient clipping magnitude. Disabled by default.\')\nflags.DEFINE_float(\'weight_decay\', 0.00004,\n                   \'Weight decay factor for training.\')\nflags.DEFINE_float(\'gpu_memory_frac_for_testing\', 0,\n                   \'If non-zero, the fraction of GPU memory that will be used. \'\n                   \'Useful for testing the benchmark script, as this allows \'\n                   \'distributed mode to be run on a single machine. For \'\n                   \'example, if there are two tasks, each can be allocated \'\n                   \'~40 percent of the memory on a single machine. This is \'\n                   \'also useful for using unified memory, as this can be set \'\n                   \'above 1 to oversubscribe the GPU using unified memory.\',\n                   lower_bound=0.)\nflags.DEFINE_boolean(\'use_unified_memory\', None,\n                     \'If True, allocate unified memory enabling larger models \'\n                     \'to fit in available device RAM.\')\nflags.DEFINE_boolean(\'timestamped_allocator\', False,\n                     \'If True marks free BFCAllocator::Chunks with time \'\n                     \'at which they are freed which can allow more efficient \'\n                     \'memory allocation in cases like RDMA networking.\')\nflags.DEFINE_integer(\'gpu_kt_max_interval\', 0,\n                     \'If > 0, the maximum number of GPU Ops that may be queued \'\n                     \'in a row without also queuing a tracking event.\')\nflags.DEFINE_integer(\'gpu_kt_max_bytes\', 0,\n                     \'If > 0, the maximum number of bytes \'\n                     \'of GPU memory that may be allocated by sequential \'\n                     \'GPU Ops without queuing a tracking event.\')\nflags.DEFINE_integer(\'gpu_kt_max_pending\', 0,\n                     \'If > 0 no more than this many GPU tracking events may be \'\n                     \'outstanding at any time.  When this limit is reached \'\n                     \'launch of additional kernels will stall until an \'\n                     \'outstanding event completes.\')\nflags.DEFINE_boolean(\'use_tf_layers\', True,\n                     \'If True, use tf.layers for neural network layers. This \'\n                     \'should not affect performance or accuracy in any way.\')\nflags.DEFINE_integer(\'tf_random_seed\', 1234,\n                     \'The TensorFlow random seed. Useful for debugging NaNs, \'\n                     \'as this can be set to various values to see if the NaNs \'\n                     \'depend on the seed.\')\nflags.DEFINE_string(\'debugger\', None,\n                    \'If set, use the TensorFlow debugger. If set to "cli", use \'\n                    \'the local CLI debugger. Otherwise, this must be in the \'\n                    \'form hostname:port (e.g., localhost:7007) in which case \'\n                    \'the experimental TensorBoard debugger will be used\')\nflags.DEFINE_boolean(\'use_python32_barrier\', False,\n                     \'When on, use threading.Barrier at Python 3.2.\')\n\nflags.DEFINE_boolean(\'ml_perf\', False,\n                     \'When True, change how the Imagenet input pipeline works \'\n                     \'slightly to meet the MLPerf compliance rules. This slows \'\n                     \'down the input pipeline. Without this option, at the end \'\n                     \'of the input pipeline, the image is divided by 127.5, \'\n                     \'then 1.0 is subtracted from it, bringing the image \'\n                     \'values from [0, 255] to [-1.0, 1.0]. With this option, \'\n                     \'each of the three channels (red, green, blue) have the \'\n                     \'average channel value among all image subtracted from \'\n                     \'it, and no division is done.\')\n\nflags.DEFINE_boolean(\'datasets_use_prefetch\', True,\n                     \'Enable use of prefetched datasets for input pipeline. \'\n                     \'This option is meaningless if use_datasets=False.\')\nflags.DEFINE_integer(\'datasets_prefetch_buffer_size\', 1,\n                     \'Prefetching op buffer size per compute device.\')\nflags.DEFINE_integer(\'datasets_num_private_threads\', None,\n                     \'Number of threads for a private threadpool created for \'\n                     \'all datasets computation. By default, we pick an \'\n                     \'appropriate number. If set to 0, we use the default \'\n                     \'tf-Compute threads for dataset operations.\')\nflags.DEFINE_boolean(\'datasets_use_caching\', False,\n                     \'Cache the compressed input data in memory. This improves \'\n                     \'the data input performance, at the cost of additional \'\n                     \'memory.\')\nflags.DEFINE_integer(\'datasets_parallel_interleave_cycle_length\', None,\n                     \'Number of parallel file readers interleaving input data.\')\nflags.DEFINE_boolean(\'datasets_sloppy_parallel_interleave\', False,\n                     \'Allow parallel interleave to depart from deterministic \'\n                     \'ordering, by temporarily skipping over files whose \'\n                     \'elements are not readily available. This can increase \'\n                     \'througput in particular in the presence of stragglers.\')\nflags.DEFINE_integer(\'datasets_parallel_interleave_prefetch\', None,\n                     \'The number of input elements to fetch before they are \'\n                     \'needed for interleaving.\')\n\nflags.DEFINE_integer(\n    \'multi_device_iterator_max_buffer_size\', 1,\n    \'Configuration parameter for the MultiDeviceIterator that \'\n    \' specifies the host side buffer size for each device.\')\n\n# Performance tuning parameters.\nflags.DEFINE_boolean(\'winograd_nonfused\', True,\n                     \'Enable/disable using the Winograd non-fused algorithms.\')\nflags.DEFINE_boolean(\n    \'batchnorm_persistent\', True,\n    \'Enable/disable using the CUDNN_BATCHNORM_SPATIAL_PERSISTENT \'\n    \'mode for batchnorm.\')\nflags.DEFINE_boolean(\'sync_on_finish\', False,\n                     \'Enable/disable whether the devices are synced after each \'\n                     \'step.\')\nflags.DEFINE_boolean(\'staged_vars\', False,\n                     \'whether the variables are staged from the main \'\n                     \'computation\')\nflags.DEFINE_boolean(\'force_gpu_compatible\', False,\n                     \'whether to enable force_gpu_compatible in GPU_Options\')\nflags.DEFINE_boolean(\'allow_growth\', None,\n                     \'whether to enable allow_growth in GPU_Options\')\nflags.DEFINE_boolean(\'xla\', False, \'whether to enable XLA auto-jit compilation\')\nflags.DEFINE_boolean(\'xla_compile\', False,\n                     \'Enable xla to compile the graph. Uncompilable ops will \'\n                     \'result in fatal errors.\')\nflags.DEFINE_boolean(\'fuse_decode_and_crop\', True,\n                     \'Fuse decode_and_crop for image preprocessing.\')\nflags.DEFINE_boolean(\'distort_color_in_yiq\', True,\n                     \'Distort color of input images in YIQ space.\')\nflags.DEFINE_boolean(\'enable_optimizations\', True,\n                     \'Whether to enable grappler and other optimizations.\')\nflags.DEFINE_string(\'rewriter_config\', None,\n                    \'Config for graph optimizers, described as a \'\n                    \'RewriterConfig proto buffer.\')\nflags.DEFINE_enum(\'loss_type_to_report\', \'total_loss\',\n                  (\'base_loss\', \'total_loss\'),\n                  \'Which type of loss to output and to write summaries for. \'\n                  \'The total loss includes L2 loss while the base loss does \'\n                  \'not. Note that the total loss is always used while \'\n                  \'computing gradients during training if weight_decay > 0, \'\n                  \'but explicitly computing the total loss, instead of just \'\n                  \'computing its gradients, can have a performance impact.\')\nflags.DEFINE_boolean(\'single_l2_loss_op\', False,\n                     \'If True, instead of using an L2 loss op per variable, \'\n                     \'concatenate the variables into a single tensor and do a \'\n                     \'single L2 loss on the concatenated tensor.\')\nflags.DEFINE_boolean(\'use_resource_vars\', False,\n                     \'Use resource variables instead of normal variables. \'\n                     \'Resource variables are slower, but this option is useful \'\n                     \'for debugging their performance.\')\nflags.DEFINE_boolean(\'compute_lr_on_cpu\', False,\n                     \'If True, do computations related to learning rate on the \'\n                     \'CPU instead of the GPU. This will significantly improve \'\n                     \'XLA performance in some cases.\')\nflags.DEFINE_boolean(\'sparse_to_dense_grads\', False,\n                     \'If True, convert all sparse gradients to dense gradients \'\n                     \'before passing them to the optimizer to update \'\n                     \'variables. Only affects models with sparse gradients, \'\n                     \'which currently is only the NCF model.\')\n# Performance tuning specific to MKL.\nflags.DEFINE_boolean(\'mkl\', False, \'If true, set MKL environment variables.\')\nflags.DEFINE_integer(\'kmp_blocktime\', 0,\n                     \'The time, in milliseconds, that a thread should wait, \'\n                     \'after completing the execution of a parallel region, \'\n                     \'before sleeping\')\nflags.DEFINE_string(\'kmp_affinity\', \'granularity=fine,verbose,compact,1,0\',\n                    \'Restricts execution of certain threads (virtual execution \'\n                    \'units) to a subset of the physical processing units in a \'\n                    \'multiprocessor computer.\')\nflags.DEFINE_integer(\'kmp_settings\', 1,\n                     \'If set to 1, MKL settings will be printed.\')\n\n# fp16 parameters. If use_fp16=False, no other fp16 parameters apply.\nflags.DEFINE_boolean(\'use_fp16\', False,\n                     \'Use 16-bit floats for certain tensors instead of 32-bit \'\n                     \'floats. This is currently experimental.\')\n# TODO(reedwm): The default loss scale of 128 causes most models to diverge\n# on the second step with synthetic data. Changing the tf.set_random_seed\n# call to tf.set_random_seed(1235) or most other seed values causes the\n# issue not to occur.\nflags.DEFINE_float(\'fp16_loss_scale\', None,\n                   \'If fp16 is enabled, the loss is multiplied by this amount \'\n                   \'right before gradients are computed, then each gradient \'\n                   \'is divided by this amount. Mathematically, this has no \'\n                   \'effect, but it helps avoid fp16 underflow. Set to 1 to \'\n                   \'effectively disable. Ignored during eval.\')\nflags.DEFINE_boolean(\'fp16_vars\', False,\n                     \'If fp16 is enabled, also use fp16 for variables. If \'\n                     \'False, the variables are stored in fp32 and casted to \'\n                     \'fp16 when retrieved.  Recommended to leave as False.\')\nflags.DEFINE_boolean(\'fp16_enable_auto_loss_scale\', False,\n                     \'If True and use_fp16 is True, automatically adjust the \'\n                     \'loss scale during training.\')\nflags.DEFINE_integer(\'fp16_inc_loss_scale_every_n\', 1000,\n                     \'If fp16 is enabled and fp16_enable_auto_loss_scale is \'\n                     \'True, increase the loss scale every n steps.\')\n\n# The method for managing variables:\n#   parameter_server: variables are stored on a parameter server that holds\n#       the master copy of the variable. In local execution, a local device\n#       acts as the parameter server for each variable; in distributed\n#       execution, the parameter servers are separate processes in the\n#       cluster.\n#       For each step, each tower gets a copy of the variables from the\n#       parameter server, and sends its gradients to the param server.\n#   replicated: each GPU has its own copy of the variables. To apply\n#       gradients, an all_reduce algorithm or or regular cross-device\n#       aggregation is used to replicate the combined gradients to all\n#       towers (depending on all_reduce_spec parameter setting).\n#   independent: each GPU has its own copy of the variables, and gradients\n#       are not shared between towers. This can be used to check performance\n#       when no data is moved between GPUs.\n#   distributed_replicated: Distributed training only. Each GPU has a copy\n#       of the variables, and updates its copy after the parameter servers\n#       are all updated with the gradients from all servers. Only works with\n#       cross_replica_sync=true. Unlike \'replicated\', currently never uses\n#       nccl all-reduce for replicating within a server.\n#   distributed_all_reduce: Distributed training where all replicas run\n#       in a single session, using all-reduce to mutally reduce the\n#       gradients.  Uses no parameter servers.  When there is only one\n#       worker, this is the same as replicated.\n#   collective_all_reduce: Distributed training where all replicas run\n#       independepently except for variable initialization and for\n#       gradient reduction which is done via collective all-reduce.\n#       NOTE: collective_all_reduce in conjunction with use_fp16 can\n#       lead to NaNs in some models (resnet50).  TODO(tucker): fix it.\n#   horovod: Distributed training using Horovod library. Runs workers using\n#       an MPI framework (e.g. Open MPI). Each worker runs training on\n#       single GPU, and averages gradients using NCCL or MPI all-reduce.\n#       See https://github.com/uber/horovod for more details.\nflags.DEFINE_enum(\'variable_update\', \'parameter_server\',\n                  (\'parameter_server\', \'replicated\', \'distributed_replicated\',\n                   \'independent\', \'distributed_all_reduce\',\n                   \'collective_all_reduce\', \'horovod\'),\n                  \'The method for managing variables: parameter_server, \'\n                  \'replicated, distributed_replicated, independent, \'\n                  \'distributed_all_reduce, collective_all_reduce, horovod\')\nflags.DEFINE_string(\'all_reduce_spec\', None,\n                    \'A specification of the all_reduce algorithm to be used \'\n                    \'for reducing gradients.  For more details, see \'\n                    \'parse_all_reduce_spec in variable_mgr.py.  An \'\n                    \'all_reduce_spec has BNF form:\\n\'\n                    \'int ::= positive whole number\\n\'\n                    \'g_int ::= int[KkMGT]?\\n\'\n                    \'alg_spec ::= alg | alg#int\\n\'\n                    \'range_spec ::= alg_spec | alg_spec/alg_spec\\n\'\n                    \'spec ::= range_spec | range_spec:g_int:range_spec\\n\'\n                    \'NOTE: not all syntactically correct constructs are \'\n                    \'supported.\\n\\n\'\n                    \'Examples:\\n \'\n                    \'"xring" == use one global ring reduction for all \'\n                    \'tensors\\n\'\n                    \'"pscpu" == use CPU at worker 0 to reduce all tensors\\n\'\n                    \'"nccl" == use NCCL to locally reduce all tensors.  \'\n                    \'Limited to 1 worker.\\n\'\n                    \'"nccl/xring" == locally (to one worker) reduce values \'\n                    \'using NCCL then ring reduce across workers.\\n\'\n                    \'"pscpu:32k:xring" == use pscpu algorithm for tensors of \'\n                    \'size up to 32kB, then xring for larger tensors.\')\n\n# If variable_update==distributed_all_reduce then it may be advantageous\n# to aggregate small tensors into one prior to reduction.  These parameters\n# control that aggregation.\nflags.DEFINE_integer(\'agg_small_grads_max_bytes\', 0,\n                     \'If > 0, try to aggregate tensors of less than this \'\n                     \'number of bytes prior to all-reduce.\')\nflags.DEFINE_integer(\'agg_small_grads_max_group\', 10,\n                     \'When aggregating small tensors for all-reduce do not \'\n                     \'aggregate more than this many into one new tensor.\')\nflags.DEFINE_integer(\'allreduce_merge_scope\', 1,\n                     \'Establish a name scope around this many \'\n                     \'gradients prior to creating the all-reduce operations. \'\n                     \'It may affect the ability of the backend to merge \'\n                     \'parallel ops.\')\n\n# Distributed training parameters.\nflags.DEFINE_enum(\'job_name\', \'\', (\'ps\', \'worker\', \'controller\', \'\'),\n                  \'One of "ps", "worker", "controller", "".  Empty for local \'\n                  \'training\')\nflags.DEFINE_string(\'ps_hosts\', \'\', \'Comma-separated list of target hosts\')\nflags.DEFINE_string(\'worker_hosts\', \'\', \'Comma-separated list of target hosts\')\nflags.DEFINE_string(\'controller_host\', None, \'optional controller host\')\nflags.DEFINE_integer(\'task_index\', 0, \'Index of task within the job\')\nflags.DEFINE_string(\'server_protocol\', \'grpc\', \'protocol for servers\')\nflags.DEFINE_boolean(\'cross_replica_sync\', True, \'\')\nflags.DEFINE_string(\'horovod_device\', \'\', \'Device to do Horovod all-reduce on: \'\n                    \'empty (default), cpu or gpu. Default with utilize GPU if \'\n                    \'Horovod was compiled with the HOROVOD_GPU_ALLREDUCE \'\n                    \'option, and CPU otherwise.\')\n\n# Summary and Save & load checkpoints.\nflags.DEFINE_integer(\'summary_verbosity\', 0, \'Verbosity level for summary ops. \'\n                     \'level 0: disable any summary.\\n\'\n                     \'level 1: small and fast ops, e.g.: learning_rate, \'\n                     \'total_loss.\\n\'\n                     \'level 2: medium-cost ops, e.g. histogram of all \'\n                     \'gradients.\\n\'\n                     \'level 3: expensive ops: images and histogram of each \'\n                     \'gradient.\\n\')\nflags.DEFINE_integer(\'save_summaries_steps\', 0,\n                     \'How often to save summaries for trained models. Pass 0 \'\n                     \'to disable summaries.\')\nflags.DEFINE_integer(\'save_model_secs\', 0,\n                     \'How often to save trained models. Pass 0 to disable \'\n                     \'saving checkpoints every N seconds. A checkpoint is \'\n                     \'saved after training completes regardless of this \'\n                     \'option.\')\nflags.DEFINE_integer(\'save_model_steps\', None,\n                     \'How often to save trained models. If specified, \'\n                     \'save_model_secs must not be specified.\')\nflags.DEFINE_integer(\'max_ckpts_to_keep\', 5,\n                     \'Max number of checkpoints to keep.\')\nflags.DEFINE_string(\'train_dir\', None,\n                    \'Path to session checkpoints. Pass None to disable saving \'\n                    \'checkpoint at the end.\')\nflags.DEFINE_string(\'eval_dir\', \'/tmp/tf_cnn_benchmarks/eval\',\n                    \'Directory where to write eval event logs.\')\nflags.DEFINE_string(\'backbone_model_path\', None,\n                    \'Path to pretrained backbone model checkpoint. Pass None \'\n                    \'if not using a backbone model.\')\nflags.DEFINE_enum(\'trt_mode\', \'\', [\'\', \'FP32\', \'FP16\', \'INT8\'],\n                  \'If this is specified in forward_only mode and \'\n                  \'freeze_when_forward_only is set to True, use TensorRT to \'\n                  \'optimize the graph before execution.\')\nflags.DEFINE_integer(\'trt_max_workspace_size_bytes\', 4 << 30,\n                     \'Max workspace size bytes used by the TensorRT optimizer.\')\n\n# Benchmark logging for model garden metric\nflags.DEFINE_string(\'benchmark_log_dir\', None,\n                    \'The directory to place the log files containing the \'\n                    \'results of benchmark. The logs are created by \'\n                    \'BenchmarkFileLogger. Requires the root of the Tensorflow \'\n                    \'models repository to be in $PYTHTONPATH.\')\nflags.DEFINE_string(\'benchmark_test_id\', None,\n                    \'The unique test ID of the benchmark run. It could be the \'\n                    \'combination of key parameters. It is hardware independent \'\n                    \'and could be used compare the performance between \'\n                    \'different test runs. This flag is designed for human \'\n                    \'consumption, and does not have any impact within the \'\n                    \'system.\')\n\nplatforms_util.define_platform_params()\n\n\nclass GlobalStepWatcher(threading.Thread):\n  """A helper class for global_step.\n\n  Polls for changes in the global_step of the model, and finishes when the\n  number of steps for the global run are done.\n  """\n\n  def __init__(self, sess, global_step_op, start_at_global_step,\n               end_at_global_step):\n    threading.Thread.__init__(self)\n    self.sess = sess\n    self.global_step_op = global_step_op\n    self.start_at_global_step = start_at_global_step\n    self.end_at_global_step = end_at_global_step\n\n    self.start_time = 0\n    self.start_step = 0\n    self.finish_time = 0\n    self.finish_step = 0\n\n  def run(self):\n    while self.finish_time == 0:\n      time.sleep(.25)\n      global_step_val, = self.sess.run([self.global_step_op])\n      if self.start_time == 0 and global_step_val >= self.start_at_global_step:\n        # Use tf.logging.info instead of log_fn, since print (which is log_fn)\n        # is not thread safe and may interleave the outputs from two parallel\n        # calls to print, which can break tests.\n        tf.logging.info(\'Starting real work at step %s at time %s\' %\n                        (global_step_val, time.ctime()))\n        self.start_time = time.perf_counter()\n        self.start_step = global_step_val\n      if self.finish_time == 0 and global_step_val >= self.end_at_global_step:\n        tf.logging.info(\'Finishing real work at step %s at time %s\' %\n                        (global_step_val, time.ctime()))\n        self.finish_time = time.perf_counter()\n        self.finish_step = global_step_val\n\n  def done(self):\n    return self.finish_time > 0\n\n  def num_steps(self):\n    return self.finish_step - self.start_step\n\n  def elapsed_time(self):\n    return self.finish_time - self.start_time\n\n\nclass CheckpointNotFoundException(Exception):\n  pass\n\n\ndef create_config_proto(params):\n  """Returns session config proto.\n\n  Args:\n    params: Params tuple, typically created by make_params or\n            make_params_from_flags.\n  """\n  config = tf.ConfigProto()\n  config.allow_soft_placement = True\n  if params.num_intra_threads is None:\n    if params.device == \'gpu\':\n      config.intra_op_parallelism_threads = 1\n  else:\n    config.intra_op_parallelism_threads = params.num_intra_threads\n  config.inter_op_parallelism_threads = params.num_inter_threads\n  config.experimental.collective_group_leader = \'/job:worker/replica:0/task:0\'\n  config.gpu_options.experimental.collective_ring_order = params.gpu_indices\n  config.gpu_options.force_gpu_compatible = params.force_gpu_compatible\n  config.experimental.use_numa_affinity = params.use_numa_affinity\n  if params.device == \'cpu\':\n    # TODO(tucker): change num_gpus to num_devices\n    config.device_count[\'CPU\'] = params.num_gpus\n  if params.allow_growth is not None:\n    config.gpu_options.allow_growth = params.allow_growth\n  if params.gpu_memory_frac_for_testing > 0:\n    config.gpu_options.per_process_gpu_memory_fraction = (\n        params.gpu_memory_frac_for_testing)\n  if params.use_unified_memory:\n    config.gpu_options.experimental.use_unified_memory = (\n        params.use_unified_memory)\n  if params.timestamped_allocator:\n    config.gpu_options.experimental.timestamped_allocator = (\n        params.timestamped_allocator)\n  if params.gpu_kt_max_interval > 0:\n    config.gpu_options.experimental.kernel_tracker_max_interval = (\n        params.gpu_kt_max_interval)\n  if params.gpu_kt_max_bytes > 0:\n    config.gpu_options.experimental.kernel_tracker_max_bytes = (\n        params.gpu_kt_max_bytes)\n  if params.gpu_kt_max_pending > 0:\n    config.gpu_options.experimental.kernel_tracker_max_pending = (\n        params.gpu_kt_max_pending)\n  if params.xla:\n    config.graph_options.optimizer_options.global_jit_level = (\n        tf.OptimizerOptions.ON_1)\n  if params.rewriter_config:\n    rewriter_config = rewriter_config_pb2.RewriterConfig()\n    text_format.Merge(params.rewriter_config, rewriter_config)\n    config.graph_options.rewrite_options.CopyFrom(rewriter_config)\n  elif not params.enable_optimizations:\n    config.graph_options.optimizer_options.opt_level = tf.OptimizerOptions.L0\n    config.graph_options.rewrite_options.disable_meta_optimizer = True\n  elif params.variable_update == \'collective_all_reduce\':\n    rewrite_options = config.graph_options.rewrite_options\n    rewrite_options.scoped_allocator_optimization = (\n        rewriter_config_pb2.RewriterConfig.ON)\n    rewrite_options.scoped_allocator_opts.enable_op.append(\'CollectiveReduce\')\n  if params.variable_update == \'horovod\':\n    import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n    config.gpu_options.visible_device_list = str(hvd.local_rank())\n  # For collective_all_reduce, ignore all devices except current worker.\n  if params.variable_update == \'collective_all_reduce\':\n    del config.device_filters[:]\n    config.device_filters.append(\n        \'/job:%s/replica:0/task:%d\' % (params.job_name, params.task_index))\n\n  # TODO(b/117324590): Re-enable PinToHostOptimizer when b/117324590 is fixed.\n  # Currently we have to disable PinToHostOptimizer w/ XLA since it causes\n  # OOM/perf cliffs.\n  config.graph_options.rewrite_options.pin_to_host_optimization = (\n      rewriter_config_pb2.RewriterConfig.OFF)\n  return config\n\n\ndef get_mode_from_params(params):\n  """Returns the mode in which this script is running.\n\n  Args:\n    params: Params tuple, typically created by make_params or\n            make_params_from_flags.\n  Raises:\n    ValueError: Unsupported params settings.\n  """\n  if params.forward_only and params.eval:\n    raise ValueError(\'Only one of forward_only and eval parameters is true\')\n\n  if params.eval:\n    return constants.BenchmarkMode.EVAL\n  elif params.forward_only:\n    return constants.BenchmarkMode.FORWARD_ONLY\n  elif (params.eval_during_training_every_n_steps or\n        params.eval_during_training_every_n_epochs or\n        params.eval_during_training_at_specified_steps or\n        params.eval_during_training_at_specified_epochs):\n    return constants.BenchmarkMode.TRAIN_AND_EVAL\n  else:\n    return constants.BenchmarkMode.TRAIN\n\n\n# How many digits to show for the loss and accuracies during training.\nLOSS_AND_ACCURACY_DIGITS_TO_SHOW = 3\n\n\ndef benchmark_one_step(sess,\n                       fetches,\n                       step,\n                       batch_size,\n                       step_train_times,\n                       trace_filename,\n                       partitioned_graph_file_prefix,\n                       profiler,\n                       image_producer,\n                       params,\n                       summary_op=None,\n                       show_images_per_sec=True,\n                       benchmark_logger=None,\n                       collective_graph_key=0,\n                       should_output_files=True):\n  """Advance one step of benchmarking."""\n  should_profile = profiler and 0 <= step < _NUM_STEPS_TO_PROFILE\n  need_options_and_metadata = (\n      should_profile or collective_graph_key > 0 or\n      ((trace_filename or partitioned_graph_file_prefix) and step == -2)\n  )\n  if need_options_and_metadata:\n    run_options = tf.RunOptions()\n    if (trace_filename and step == -2) or should_profile:\n      run_options.trace_level = tf.RunOptions.FULL_TRACE\n    if partitioned_graph_file_prefix and step == -2:\n      run_options.output_partition_graphs = True\n    if collective_graph_key > 0:\n      run_options.experimental.collective_graph_key = collective_graph_key\n    run_metadata = tf.RunMetadata()\n  else:\n    run_options = None\n    run_metadata = None\n  summary_str = None\n  start_time = time.perf_counter()\n  if summary_op is None:\n    results = sess.run(fetches, options=run_options, run_metadata=run_metadata)\n  else:\n    (results, summary_str) = sess.run(\n        [fetches, summary_op], options=run_options, run_metadata=run_metadata)\n\n  if not params.forward_only:\n    lossval = results[\'average_loss\']\n  else:\n    lossval = 0.\n  if image_producer is not None:\n    image_producer.notify_image_consumption()\n  train_time = time.perf_counter() - start_time\n  step_train_times.append(train_time)\n  if (show_images_per_sec and step >= 0 and\n      (step == 0 or (step + 1) % params.display_every == 0)):\n    speed_mean, speed_uncertainty, speed_jitter = get_perf_timing(\n        batch_size, step_train_times, params.display_perf_ewma)\n    log_str = \'%i\\t%s\\t%.*f\' % (\n        step + 1,\n        get_perf_timing_str(speed_mean, speed_uncertainty, speed_jitter),\n        LOSS_AND_ACCURACY_DIGITS_TO_SHOW, lossval)\n    if \'top_1_accuracy\' in results:\n      log_str += \'\\t%.*f\\t%.*f\' % (\n          LOSS_AND_ACCURACY_DIGITS_TO_SHOW, results[\'top_1_accuracy\'],\n          LOSS_AND_ACCURACY_DIGITS_TO_SHOW, results[\'top_5_accuracy\'])\n    log_fn(log_str)\n    if benchmark_logger:\n      benchmark_logger.log_metric(\n          \'current_examples_per_sec\', speed_mean, global_step=step + 1)\n      if \'top_1_accuracy\' in results:\n        benchmark_logger.log_metric(\n            \'top_1_accuracy\', results[\'top_1_accuracy\'], global_step=step + 1)\n        benchmark_logger.log_metric(\n            \'top_5_accuracy\', results[\'top_5_accuracy\'], global_step=step + 1)\n  if need_options_and_metadata:\n    if should_profile:\n      profiler.add_step(step, run_metadata)\n    if trace_filename and step == -2 and should_output_files:\n      log_fn(\'Dumping trace to %s\' % trace_filename)\n      trace_dir = os.path.dirname(trace_filename)\n      if not gfile.Exists(trace_dir):\n        gfile.MakeDirs(trace_dir)\n      with gfile.Open(trace_filename, \'w\') as trace_file:\n        if params.use_chrome_trace_format:\n          trace = timeline.Timeline(step_stats=run_metadata.step_stats)\n          trace_file.write(trace.generate_chrome_trace_format(show_memory=True))\n        else:\n          trace_file.write(str(run_metadata.step_stats))\n    if partitioned_graph_file_prefix and step == -2 and should_output_files:\n      path, filename = os.path.split(partitioned_graph_file_prefix)\n      if \'.\' in filename:\n        base_filename, ext = filename.rsplit(\'.\', 1)\n        ext = \'.\' + ext\n      else:\n        base_filename, ext = filename, \'\'\n      as_text = filename.endswith(\'txt\')\n      for graph_def in run_metadata.partition_graphs:\n        device = graph_def.node[0].device.replace(\'/\', \'_\').replace(\':\', \'_\')\n        graph_filename = \'%s%s%s\' % (base_filename, device, ext)\n        log_fn(\'Writing partitioned GraphDef as %s to %s\' % (\n            \'text\' if as_text else \'binary\',\n            os.path.join(path, graph_filename)))\n        tf.train.write_graph(graph_def, path, graph_filename, as_text)\n  return (summary_str, lossval)\n\n\ndef get_perf_timing_str(speed_mean, speed_uncertainty, speed_jitter, scale=1):\n  if scale == 1:\n    # TODO(laigd): rename \'images\' to maybe \'inputs\', same below.\n    return (\'images/sec: %.1f +/- %.1f (jitter = %.1f)\' %\n            (speed_mean, speed_uncertainty, speed_jitter))\n  else:\n    return \'images/sec: %.1f\' % speed_mean\n\n\ndef get_perf_timing(batch_size, step_train_times, ewma_alpha=None, scale=1):\n  """Calculate benchmark processing speed."""\n  times = np.array(step_train_times)\n  speeds = batch_size / times\n  if ewma_alpha:\n    weights = np.logspace(len(times)-1, 0, len(times), base=1-ewma_alpha)\n    time_mean = np.average(times, weights=weights)\n  else:\n    time_mean = np.mean(times)\n  speed_mean = scale * batch_size / time_mean\n  speed_uncertainty = np.std(speeds) / np.sqrt(float(len(speeds)))\n  speed_jitter = 1.4826 * np.median(np.abs(speeds - np.median(speeds)))\n  return speed_mean, speed_uncertainty, speed_jitter\n\n\ndef load_checkpoint(saver, sess, ckpt_dir):\n  """Loads checkpoint from provided directory or full path.\n\n  Args:\n    saver: Saver used to restore the checkpoint.\n    sess: TensorFlow session.\n    ckpt_dir: Path to a folder of checkpoints or full path to a checkpoint.\n\n  Returns:\n    Global step.\n  """\n  model_checkpoint_path = _get_checkpoint_to_load(ckpt_dir)\n  global_step = model_checkpoint_path.split(\'/\')[-1].split(\'-\')[-1]\n  if not global_step.isdigit():\n    global_step = 0\n  else:\n    global_step = int(global_step)\n  saver.restore(sess, model_checkpoint_path)\n  log_fn(\'Successfully loaded model from %s.\' % model_checkpoint_path)\n  return global_step\n\n\ndef _get_checkpoint_to_load(ckpt_dir):\n  """Returns which checkpoint to load.\n\n  Args:\n    ckpt_dir: Path to a folder of checkpoints or full path to a checkpoint.\n\n  Returns:\n    Full path to checkpoint to load.\n\n  Raises:\n    CheckpointNotFoundException: If checkpoint is not found.\n  """\n  p = re.compile(r\'ckpt-\\d+$\')\n  if p.search(ckpt_dir):\n    model_checkpoint_path = ckpt_dir\n  else:\n    # Finds latest checkpoint in directory provided\n    ckpt = tf.train.get_checkpoint_state(ckpt_dir)\n    if ckpt and ckpt.model_checkpoint_path:\n      model_checkpoint_path = ckpt.model_checkpoint_path\n    else:\n      raise CheckpointNotFoundException(\'No checkpoint file found in dir:{}\'.\n                                        format(ckpt_dir))\n  return model_checkpoint_path\n\n\n# Params are passed to BenchmarkCNN\'s constructor. Params is a map from name\n# to value, with one field per key in flags.param_specs.\n#\n# Call make_params() or make_params_from_flags() below to construct a Params\n# tuple with default values from flags.param_specs, rather than constructing\n# Params directly.\nParams = namedtuple(\'Params\', flags.param_specs.keys())  # pylint: disable=invalid-name\n\n\ndef validate_params(params):\n  """Validates that the Params tuple had valid values.\n\n  When command-line flags are defined for each ParamSpec by calling\n  flags.define_flags(), calling this function is unnecessary because absl\n  already does flag validation. Otherwise, this function should be called.\n\n  Args:\n     params: A Params tuple.\n  Raises:\n    ValueError: An element of params had an invalid value.\n  """\n  for name, value in params._asdict().items():\n    param_spec = flags.param_specs[name]\n    if param_spec.flag_type in (\'integer\', \'float\'):\n      if (value is not None and param_spec.kwargs[\'lower_bound\'] is not None and\n          value < param_spec.kwargs[\'lower_bound\']):\n        raise ValueError(\'Param %s value of %s is lower than the lower bound \'\n                         \'of %s\' %\n                         (name, value, param_spec.kwargs[\'lower_bound\']))\n      if (value is not None and param_spec.kwargs[\'upper_bound\'] is not None and\n          param_spec.kwargs[\'upper_bound\'] < value):\n        raise ValueError(\'Param %s value of %s is higher than the upper bound \'\n                         \'of %s\' %\n                         (name, value, param_spec.kwargs[\'upper_bound\']))\n    elif (value is not None and param_spec.flag_type == \'enum\' and\n          value not in param_spec.kwargs[\'enum_values\']):\n      raise ValueError(\'Param %s of value %s is not in %s\'%\n                       (name, value, param_spec.kwargs[\'enum_values\']))\n\n\ndef make_params(**kwargs):\n  """Create a Params tuple for BenchmarkCNN from kwargs.\n\n  Default values are filled in from flags.param_specs.\n\n  Args:\n    **kwargs: kwarg values will override the default values.\n  Returns:\n    Params namedtuple for constructing BenchmarkCNN.\n  """\n  # Create a (name: default_value) map from flags.param_specs.\n  default_kwargs = {\n      name: flags.param_specs[name].default_value\n      for name in flags.param_specs\n  }\n  params = Params(**default_kwargs)._replace(**kwargs)\n  validate_params(params)\n  return params\n\n\ndef make_params_from_flags():\n  """Create a Params tuple for BenchmarkCNN from absl_flags.FLAGS.\n\n  Returns:\n    Params namedtuple for constructing BenchmarkCNN.\n  """\n  # Collect (name: value) pairs for absl_flags.FLAGS with matching names in\n  # flags.param_specs.\n  flag_values = {name: getattr(absl_flags.FLAGS, name)\n                 for name in flags.param_specs.keys()}\n  return Params(**flag_values)\n\n\ndef remove_param_fields(params, fields_to_remove):\n  """Remove fields from a Params namedtuple."""\n  params_dict = params._asdict()\n  for field in fields_to_remove:\n    assert field in params_dict, \'Invalid Params field: \' + field\n  params_dict = {k: v for k, v in params_dict.items()\n                 if k not in fields_to_remove}\n  new_params_type = namedtuple(\'Params\', params_dict.keys())\n  return new_params_type(**params_dict)\n\n\ndef get_num_batches_and_epochs(params, batch_size, num_examples_per_epoch):\n  """Returns the number of batches and epochs to run for.\n\n  Args:\n    params: Params tuple, typically created by make_params or\n      make_params_from_flags.\n    batch_size: The number of images per step.\n    num_examples_per_epoch: The number of images in a single epoch.\n\n  Returns:\n    num_batches: The number of batches to run for.\n    num_epochs: The number of epochs to run for. This might be slightly\n      smaller than params.num_epochs if specified, because the number of batches\n      must be an integer.\n\n  Raises:\n    ValueError: Invalid or unsupported params.\n  """\n  if params.num_batches and params.num_epochs:\n    raise ValueError(\'At most one of --num_batches and --num_epochs may be \'\n                     \'specified.\')\n  if params.num_epochs:\n    num_batches = int(params.num_epochs * num_examples_per_epoch +\n                      batch_size - 1) // batch_size\n  else:\n    num_batches = params.num_batches or _DEFAULT_NUM_BATCHES\n  num_epochs = num_batches * batch_size / num_examples_per_epoch\n  return (num_batches, num_epochs)\n\n\ndef get_piecewise_learning_rate(piecewise_learning_rate_schedule,\n                                global_step, num_batches_per_epoch):\n  """Returns a piecewise learning rate tensor.\n\n  Args:\n    piecewise_learning_rate_schedule: The --piecewise_learning_rate_schedule\n      parameter\n    global_step: Scalar tensor representing the global step.\n    num_batches_per_epoch: float indicating the number of batches per epoch.\n\n  Returns:\n    A scalar float tensor, representing the learning rate.\n\n  Raises:\n    ValueError: piecewise_learning_rate_schedule is not formatted correctly.\n  """\n  pieces = piecewise_learning_rate_schedule.split(\';\')\n  if len(pieces) % 2 == 0:\n    raise ValueError(\'--piecewise_learning_rate_schedule must have an odd \'\n                     \'number of components\')\n  values = []\n  boundaries = []\n  for i, piece in enumerate(pieces):\n    if i % 2 == 0:\n      try:\n        values.append(float(piece))\n      except ValueError:\n        raise ValueError(\'Invalid learning rate: \' + piece)\n    else:\n      try:\n        boundaries.append(int(int(piece) * num_batches_per_epoch) - 1)\n      except ValueError:\n        raise ValueError(\'Invalid epoch: \' + piece)\n  return tf.train.piecewise_constant(global_step, boundaries, values,\n                                     name=\'piecewise_learning_rate\')\n\n\ndef get_learning_rate(params, global_step, num_examples_per_epoch, model,\n                      batch_size):\n  """Returns a learning rate tensor based on global_step.\n\n  Args:\n    params: Params tuple, typically created by make_params or\n      make_params_from_flags.\n    global_step: Scalar tensor representing the global step.\n    num_examples_per_epoch: The number of examples per epoch.\n    model: The model.Model object to obtain the default learning rate from if no\n      learning rate is specified.\n    batch_size: Number of examples per step\n\n  Returns:\n    A scalar float tensor, representing the learning rate. When evaluated, the\n    learning rate depends on the current value of global_step.\n\n  Raises:\n    ValueError: Invalid or unsupported params.\n  """\n  with tf.name_scope(\'learning_rate\'):\n    num_batches_per_epoch = num_examples_per_epoch / batch_size\n\n    if params.piecewise_learning_rate_schedule:\n      if (params.init_learning_rate is not None or\n          params.learning_rate_decay_factor or\n          params.minimum_learning_rate or params.num_epochs_per_decay):\n        raise ValueError(\'No other learning rate-related flags can be \'\n                         \'specified if --piecewise_learning_rate_schedule is \'\n                         \'specified\')\n      learning_rate = get_piecewise_learning_rate(\n          params.piecewise_learning_rate_schedule,\n          global_step, num_batches_per_epoch)\n    elif params.init_learning_rate is not None:\n      learning_rate = params.init_learning_rate\n      if (params.num_epochs_per_decay > 0 and\n          params.learning_rate_decay_factor > 0):\n        decay_steps = int(num_batches_per_epoch * params.num_epochs_per_decay)\n\n        # Decay the learning rate exponentially based on the number of steps.\n        learning_rate = tf.train.exponential_decay(\n            params.init_learning_rate,\n            global_step,\n            decay_steps,\n            params.learning_rate_decay_factor,\n            staircase=True)\n\n        if params.minimum_learning_rate != 0.:\n          learning_rate = tf.maximum(learning_rate,\n                                     params.minimum_learning_rate)\n    else:\n      learning_rate = model.get_learning_rate(global_step, batch_size)\n    if params.num_learning_rate_warmup_epochs > 0 and (\n        params.init_learning_rate is not None or\n        params.piecewise_learning_rate_schedule):\n      warmup_steps = int(num_batches_per_epoch *\n                         params.num_learning_rate_warmup_epochs)\n      init_lr = params.init_learning_rate\n      if init_lr is None:\n        init_lr = float(params.piecewise_learning_rate_schedule.split(\';\')[0])\n      warmup_lr = init_lr * tf.cast(global_step, tf.float32) / tf.cast(\n          warmup_steps, tf.float32)\n      learning_rate = tf.cond(global_step < warmup_steps,\n                              lambda: warmup_lr, lambda: learning_rate)\n\n    learning_rate = mlperf.logger.log_deferred_tensor_value(\n        mlperf.tags.OPT_LR, learning_rate, global_step, every_n=100)\n  return learning_rate\n\n\ndef get_optimizer(params, learning_rate):\n  """Returns the optimizer that should be used based on params."""\n  if params.optimizer == \'momentum\':\n    mlperf.logger.log(key=mlperf.tags.OPT_NAME,\n                      value=mlperf.tags.SGD_WITH_MOMENTUM)\n    mlperf.logger.log(key=mlperf.tags.OPT_MOMENTUM, value=params.momentum)\n    opt = tf.train.MomentumOptimizer(\n        learning_rate, params.momentum, use_nesterov=True)\n  elif params.optimizer == \'sgd\':\n    mlperf.logger.log(key=mlperf.tags.OPT_NAME, value=mlperf.tags.SGD)\n    opt = tf.train.GradientDescentOptimizer(learning_rate)\n  elif params.optimizer == \'rmsprop\':\n    opt = tf.train.RMSPropOptimizer(\n        learning_rate,\n        params.rmsprop_decay,\n        momentum=params.rmsprop_momentum,\n        epsilon=params.rmsprop_epsilon)\n  elif params.optimizer == \'adam\':\n    opt = tf.train.AdamOptimizer(learning_rate, params.adam_beta1,\n                                 params.adam_beta2, params.adam_epsilon)\n  else:\n    raise ValueError(\'Optimizer "{}" was not recognized\'.\n                     format(params.optimizer))\n  return opt\n\n\ndef generate_tfprof_profile(profiler, tfprof_file):\n  """Generates a tfprof profile, writing it to a file and printing top ops.\n\n  Args:\n    profiler: A tf.profiler.Profiler. `profiler.add_step` must have already been\n      called.\n    tfprof_file: The filename to write the ProfileProto to.\n  """\n  profile_proto = profiler.serialize_to_string()\n  log_fn(\'Dumping ProfileProto to %s\' % tfprof_file)\n  with gfile.Open(tfprof_file, \'wb\') as f:\n    f.write(profile_proto)\n\n  # Print out the execution times of the top operations. Note this\n  # information can also be obtained with the dumped ProfileProto, but\n  # printing it means tfprof doesn\'t have to be used if all the user wants\n  # is the top ops.\n  options = tf.profiler.ProfileOptionBuilder.time_and_memory()\n  options[\'max_depth\'] = _NUM_OPS_TO_PRINT\n  options[\'order_by\'] = \'accelerator_micros\'\n  profiler.profile_operations(options)\n\n\nclass BenchmarkCNN(object):\n  """Class for benchmarking a cnn network."""\n\n  def __init__(self, params, dataset=None, model=None):\n    """Initialize BenchmarkCNN.\n\n    Args:\n      params: Params tuple, typically created by make_params or\n              make_params_from_flags.\n      dataset: If not None, the dataset to use. Otherwise, params is used to\n               obtain the dataset.\n      model: If not None, the model to use. Otherwise, params is used to obtain\n             the model.\n    Raises:\n      ValueError: Unsupported params settings.\n    """\n    mlperf.logger.log(key=mlperf.tags.RUN_START)\n    self.params = params\n    if params.eval:\n      self._doing_eval = True\n    else:\n      # Note self._doing_eval can later switch to True in self._do_eval() if\n      # self.params.eval_during_training_* is specified.\n      self._doing_eval = False\n    self.dataset = dataset or datasets.create_dataset(self.params.data_dir,\n                                                      self.params.data_name)\n    self.model = model or model_config.get_model_config(\n        self.params.model, self.dataset, self.params)\n    self.trace_filename = self.params.trace_file\n    self.rewriter_config = self.params.rewriter_config\n    autotune_threshold = self.params.autotune_threshold if (\n        self.params.autotune_threshold) else 1\n    min_autotune_warmup = 5 * autotune_threshold * autotune_threshold\n    self.num_warmup_batches = self.params.num_warmup_batches if (\n        self.params.num_warmup_batches is not None) else max(\n            10, min_autotune_warmup)\n    self.graph_file = self.params.graph_file\n    self.resize_method = self.params.resize_method\n    self.sync_queue_counter = 0\n    self.num_gpus = self.params.num_gpus\n    if self.params.gpu_indices:\n      self.gpu_indices = [int(x) for x in self.params.gpu_indices.split(\',\')]\n    else:\n      self.gpu_indices = [x for x in range(self.num_gpus)]\n\n    if (self.params.device == \'cpu\' and self.params.data_format == \'NCHW\' and\n        not self.params.mkl):\n      raise ValueError(\'device=cpu requires that data_format=NHWC\')\n\n    if ((self.params.num_epochs_per_decay or\n         self.params.learning_rate_decay_factor) and\n        not (self.params.init_learning_rate is not None and\n             self.params.num_epochs_per_decay\n             and self.params.learning_rate_decay_factor)):\n      raise ValueError(\'If one of num_epochs_per_decay or \'\n                       \'learning_rate_decay_factor is set, both must be set\'\n                       \'and learning_rate must be set\')\n    if (self.params.minimum_learning_rate and\n        not (self.params.init_learning_rate is not None and\n             self.params.num_epochs_per_decay and\n             self.params.learning_rate_decay_factor)):\n      raise ValueError(\'minimum_learning_rate requires learning_rate,\'\n                       \'num_epochs_per_decay, and \'\n                       \'learning_rate_decay_factor to be set\')\n\n    if (self.params.use_fp16 and self.params.fp16_vars and\n        \'replicated\' in self.params.variable_update and\n        self.params.all_reduce_spec and \'nccl\' in self.params.all_reduce_spec):\n      raise ValueError(\'fp16 variables are not supported with NCCL\')\n    if (self.params.use_fp16 and self.params.fp16_vars and\n        self.params.gradient_repacking):\n      raise ValueError(\'--fp16_vars cannot be used with --gradient_repacking\')\n\n    if self.params.variable_update == \'horovod\' and self.params.num_gpus > 1:\n      raise ValueError(\'Horovod benchmarks require num_gpus=1 on each worker\')\n\n    if self.params.variable_update == \'horovod\' and self.params.job_name:\n      raise ValueError(\'job_name should not be specified for Horovod.\')\n\n    if self.params.use_fp16 and self.params.fp16_enable_auto_loss_scale:\n      if self.params.all_reduce_spec and \'nccl\' in self.params.all_reduce_spec:\n        raise ValueError(\'Automatic loss scaling is not supported with NCCL.\')\n      if self.params.variable_update not in (\'parameter_server\', \'replicated\',\n                                             \'independent\'):\n        raise ValueError(\'Automatic loss scaling is not supported with \'\n                         \'variable_update=%s.\' % self.params.variable_update)\n      if self.params.staged_vars:\n        raise ValueError(\'Automatic loss scaling is not supported with\'\n                         \'staged_vars.\')\n\n    if (self.params.debugger is not None and self.params.debugger != \'cli\' and\n        \':\' not in self.params.debugger):\n      raise ValueError(\'--debugger must be "cli" or in the form \'\n                       \'host:port\')\n\n    if self.params.hierarchical_copy and self.params.num_gpus <= 1:\n      raise ValueError(\'--hierarchical_copy requires --num_gpus to be greater \'\n                       \'than 1\')\n\n    if params.save_model_secs and params.save_model_steps:\n      raise ValueError(\'At most one of --save_model_secs and \'\n                       \'--save_model_steps can be specified\')\n\n    eval_during_training_flags = list(map(bool, [\n        params.eval_during_training_every_n_steps,\n        params.eval_during_training_every_n_epochs,\n        params.eval_during_training_at_specified_steps,\n        params.eval_during_training_at_specified_epochs,\n    ]))\n\n    if eval_during_training_flags.count(True) > 1:\n      raise ValueError(\'At most one flag with --eval_during_training_* prefix \'\n                       \'must be specified.\')\n\n    eval_during_training_enabled = any(eval_during_training_flags)\n\n    if eval_during_training_enabled:\n      if params.eval:\n        raise ValueError(\'At most one of --eval and --eval_during_training_* \'\n                         \'must be specified\')\n      if params.forward_only:\n        raise ValueError(\'At most one of --forward_only and \'\n                         \'--eval_during_training_* must be specified\')\n      if params.job_name:\n        raise ValueError(\'--eval_during_training_* is not yet supported in \'\n                         \'distributed mode.\')\n      if params.staged_vars:\n        raise ValueError(\'--eval_during_training_* is not currently compatible \'\n                         \'with --staged_vars\')\n\n    if params.stop_at_top_1_accuracy and not eval_during_training_enabled:\n      raise ValueError(\'--stop_at_top_1_accuracy is only supported with \'\n                       \'--eval_during_training_*\')\n    if params.collect_eval_results_async and params.model != \'ssd300\':\n      raise ValueError(\'--collect_eval_results_async only works with ssd300 \'\n                       \'model currently.\')\n    if self.params.forward_only and self.params.freeze_when_forward_only:\n      if self.params.train_dir is not None:\n        raise ValueError(\'In forward_only mode, when --freeze_when_forward_only\'\n                         \' is True, --train_dir should not be specified\')\n      if self.params.data_dir and not self.params.datasets_use_prefetch:\n        raise ValueError(\'In forward_only mode, when --freeze_when_forward_only\'\n                         \' is True and --data_dir is set, \'\n                         \'--datasets_use_prefetch should be set to True\')\n      if self.params.job_name:\n        raise ValueError(\'In forward_only mode, when --freeze_when_forward_only\'\n                         \' is True, --job_name should not be specified and \'\n                         \'distributed running is not supported\')\n      self.forward_only_and_freeze = True\n    else:\n      self.forward_only_and_freeze = False\n      if self.params.trt_mode:\n        raise ValueError(\'--trt_mode should not be specified if one of \'\n                         \'--forward_only and --freeze_when_forward_only is set \'\n                         \'to False\')\n\n    self.mode = get_mode_from_params(self.params)\n\n    # Use the batch size from the command line if specified, otherwise use the\n    # model\'s default batch size.  Scale the benchmark\'s batch size by the\n    # number of GPUs.\n    if self.params.batch_size > 0:\n      self.model.set_batch_size(self.params.batch_size)\n    self.batch_size = self.model.get_batch_size() * self.num_gpus\n    if self.mode in (constants.BenchmarkMode.TRAIN,\n                     constants.BenchmarkMode.TRAIN_AND_EVAL):\n      self.train_batch_size = self.batch_size\n    else:\n      self.train_batch_size = None\n    if self.mode in (constants.BenchmarkMode.EVAL,\n                     constants.BenchmarkMode.TRAIN_AND_EVAL):\n      if self.params.eval_batch_size > 0:\n        self.eval_batch_size = self.params.eval_batch_size * self.num_gpus\n      else:\n        self.eval_batch_size = self.batch_size\n    else:\n      self.eval_batch_size = None\n    self.batch_group_size = self.params.batch_group_size\n    self.enable_auto_loss_scale = (\n        self.params.use_fp16 and self.params.fp16_enable_auto_loss_scale)\n    self.loss_scale = None\n    self.loss_scale_normal_steps = None\n\n    self.job_name = self.params.job_name  # "" for local training\n\n    # PS server is used for distributed jobs not using all-reduce.\n    use_ps_server = self.job_name and (self.params.variable_update !=\n                                       \'distributed_all_reduce\' and\n                                       self.params.variable_update !=\n                                       \'collective_all_reduce\')\n    # controller is used for distributed_all_reduce with > 1 worker.\n    use_controller = (\n        self.params.variable_update == \'distributed_all_reduce\' and\n        self.job_name)\n    if use_controller and not params.controller_host:\n      raise ValueError(\'When variable_update==distributed_all_reduce \'\n                       \'controller_host must also be specified.\')\n    self.single_session = (\n        self.params.variable_update == \'distributed_all_reduce\')\n    # collective_all_reduce doesn\'t need a controller or ps\n    self.distributed_collective = (\n        self.params.variable_update == \'collective_all_reduce\' and\n        self.job_name)\n\n    self.local_parameter_device_flag = self.params.local_parameter_device\n    if self.job_name:\n      self.task_index = self.params.task_index\n      self.cluster_manager = platforms_util.get_cluster_manager(\n          params, create_config_proto(params))\n      assert isinstance(self.cluster_manager, cnn_util.BaseClusterManager)\n\n      worker_prefix = \'/job:worker/replica:0/task:%s\' % self.task_index\n      if use_ps_server:\n        self.param_server_device = tf.train.replica_device_setter(\n            worker_device=worker_prefix + \'/cpu:0\',\n            cluster=self.cluster_manager.get_cluster_spec())\n        # This device on which the queues for managing synchronization between\n        # servers should be stored.\n        self.sync_queue_devices = [\n            \'/job:ps/replica:0/task:%s/cpu:0\' % i\n            for i in range(self.cluster_manager.num_ps())\n        ]\n      else:\n        self.sync_queue_devices = [\'/job:worker/replica:0/task:0/cpu:0\']\n    else:\n      self.task_index = 0\n      self.cluster_manager = None\n      worker_prefix = \'\'\n      self.param_server_device = \'/%s:0\' % self.params.local_parameter_device\n      self.sync_queue_devices = [self.param_server_device]\n\n    if self.cluster_manager:\n      self.num_workers = self.cluster_manager.num_workers()\n    elif self.params.variable_update == \'horovod\':\n      import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n      self.num_workers = hvd.size()\n    else:\n      self.num_workers = 1\n    self.num_ps = self.cluster_manager.num_ps() if self.cluster_manager else 0\n\n    if self.num_workers > 1 and self.params.all_reduce_spec == \'nccl\':\n      raise ValueError(\'--all_reduce_spec=nccl is invalid in a \'\n                       \'multi-worker job\')\n\n    # Device to use for ops that need to always run on the local worker\'s CPU.\n    self.cpu_device = \'%s/cpu:0\' % worker_prefix\n\n    # Device to use for ops that need to always run on the local worker\'s\n    # compute device, and never on a parameter server device.\n    self.raw_devices = [\n        \'%s/%s:%i\' % (worker_prefix, self.params.device, i)\n        for i in xrange(self.num_gpus)\n    ]\n\n    subset = \'validation\' if params.eval else \'train\'\n    self.num_batches, self.num_epochs = get_num_batches_and_epochs(\n        params, self.batch_size * self.num_workers,\n        self.dataset.num_examples_per_epoch(subset))\n    if self.mode in (constants.BenchmarkMode.EVAL,\n                     constants.BenchmarkMode.TRAIN_AND_EVAL):\n      # TODO(reedwm): Currently we do extra eval logic for num_eval_batches and\n      # the preprocessor. We should encapsulate this logic into a shared\n      # function or class.\n      if params.num_eval_batches is None and params.num_eval_epochs is None:\n        eval_params = self.params\n      else:\n        eval_params = self.params._replace(\n            num_batches=self.params.num_eval_batches,\n            num_epochs=self.params.num_eval_epochs)\n      self.num_eval_batches, self.num_eval_epochs = get_num_batches_and_epochs(\n          eval_params, self.eval_batch_size * self.num_workers,\n          self.dataset.num_examples_per_epoch(\'validation\'))\n    else:\n      self.num_eval_batches, self.num_eval_epochs = None, None\n\n    num_train_examples_per_epoch = self.dataset.num_examples_per_epoch(\'train\')\n    if self.params.eval_during_training_every_n_epochs:\n      n_epochs = self.params.eval_during_training_every_n_epochs\n      self.eval_during_training_at_specified_steps = {\n          (int(e * num_train_examples_per_epoch + self.batch_size - 1) //\n           self.batch_size)\n          for e in np.arange(n_epochs, self.num_epochs, n_epochs)}\n\n    if self.params.eval_during_training_at_specified_steps:\n      try:\n        self.eval_during_training_at_specified_steps = set(map(\n            int, self.params.eval_during_training_at_specified_steps))\n      except ValueError:\n        raise ValueError(\'Param eval_during_training_at_specified_steps value \'\n                         \'of %s cannot be converted to a list of integers.\' %\n                         (self.params.eval_during_training_at_specified_steps))\n\n    if self.params.eval_during_training_at_specified_epochs:\n      try:\n        n_epochs = list(map(\n            float, self.params.eval_during_training_at_specified_epochs))\n        offset = n_epochs[0] - 1\n        if offset.is_integer():\n          offset = int(offset)\n        mlperf.logger.log(key=mlperf.tags.EVAL_EPOCH_OFFSET, value=offset)\n        self.eval_during_training_at_specified_steps = {\n            (int(e * num_train_examples_per_epoch + self.batch_size - 1) //\n             self.batch_size)\n            for e in n_epochs}\n      except ValueError:\n        raise ValueError(\'Param eval_during_training_at_specified_epochs value \'\n                         \'of %s cannot be converted to a list of floats.\' %\n                         (self.params.eval_during_training_at_specified_epochs))\n\n    if params.eval_during_training_every_n_epochs:\n      offset = params.eval_during_training_every_n_epochs - 1\n      if offset.is_integer():\n        offset = int(offset)\n      mlperf.logger.log(key=mlperf.tags.EVAL_EPOCH_OFFSET, value=offset)\n\n    if (self.params.staged_vars and\n        self.params.variable_update != \'parameter_server\'):\n      raise ValueError(\'staged_vars for now is only supported with \'\n                       \'variable_update=parameter_server\')\n\n    if self.params.variable_update == \'parameter_server\':\n      if self.job_name:\n        if not self.params.staged_vars:\n          self.variable_mgr = variable_mgr.VariableMgrDistributedFetchFromPS(\n              self)\n        else:\n          self.variable_mgr = (\n              variable_mgr.VariableMgrDistributedFetchFromStagedPS(self))\n      else:\n        if not self.params.staged_vars:\n          self.variable_mgr = variable_mgr.VariableMgrLocalFetchFromPS(self)\n        else:\n          self.variable_mgr = variable_mgr.VariableMgrLocalFetchFromStagedPS(\n              self)\n    elif self.params.variable_update == \'replicated\':\n      if self.job_name:\n        raise ValueError(\'Invalid variable_update in distributed mode: %s\' %\n                         self.params.variable_update)\n      self.variable_mgr = variable_mgr.VariableMgrLocalReplicated(\n          self, self.params.all_reduce_spec,\n          self.params.agg_small_grads_max_bytes,\n          self.params.agg_small_grads_max_group,\n          self.params.allreduce_merge_scope)\n    elif self.params.variable_update == \'distributed_all_reduce\':\n      assert self.params.cross_replica_sync\n      self.variable_mgr = variable_mgr.VariableMgrDistributedAllReduce(\n          self, self.params.all_reduce_spec,\n          (\'worker\' if self.num_workers > 1 else \'localhost\'),\n          self.num_workers, self.params.agg_small_grads_max_bytes,\n          self.params.agg_small_grads_max_group,\n          self.params.allreduce_merge_scope)\n    elif self.params.variable_update == \'collective_all_reduce\':\n      assert self.params.cross_replica_sync\n      self.variable_mgr = variable_mgr.VariableMgrCollectiveAllReduce(\n          self, self.params.all_reduce_spec,\n          self.num_workers, self.num_gpus, self.task_index,\n          self.params.allreduce_merge_scope)\n    elif self.params.variable_update == \'distributed_replicated\':\n      assert self.params.cross_replica_sync\n      if not self.job_name:\n        raise ValueError(\'Invalid variable_update in local mode: %s\' %\n                         self.params.variable_update)\n      self.variable_mgr = variable_mgr.VariableMgrDistributedReplicated(self)\n    elif self.params.variable_update in (\'independent\', \'horovod\'):\n      if self.job_name:\n        raise ValueError(\'Invalid variable_update in distributed mode: %s\' %\n                         self.params.variable_update)\n      self.variable_mgr = variable_mgr.VariableMgrIndependent(self)\n    else:\n      raise ValueError(\n          \'Invalid variable_update: %s\' % self.params.variable_update)\n\n    # Device to use for running on the local worker\'s compute device, but\n    # with variables assigned to parameter server devices.\n    self.devices = self.variable_mgr.get_devices()\n    if self.job_name:\n      if use_ps_server:\n        self.global_step_device = self.param_server_device\n      elif self.params.variable_update == \'collective_all_reduce\':\n        self.global_step_device = self.cpu_device\n      else:\n        self.global_step_device = \'/job:worker/replica:0/task:0/cpu:0\'\n    else:\n      self.global_step_device = self.cpu_device\n\n    self.input_preprocessor = None\n    self.eval_input_preprocessor = None\n    if not self.dataset.use_synthetic_gpu_inputs():\n      if not self.params.eval:\n        self.input_preprocessor = self.get_input_preprocessor()\n      if self.mode in (constants.BenchmarkMode.EVAL,\n                       constants.BenchmarkMode.TRAIN_AND_EVAL):\n        with self._do_eval():\n          self.eval_input_preprocessor = self.get_input_preprocessor()\n    self.datasets_use_prefetch = (\n        self.params.datasets_use_prefetch and\n        # TODO(rohanj): Figure out why --datasets_use_prefetch freezes on the\n        # CPU.\n        self.params.device.lower() != \'cpu\' and\n        self.input_preprocessor and\n        self.input_preprocessor.supports_datasets())\n    self.init_global_step = 0\n\n    self._config_benchmark_logger()\n\n    if self.mode == constants.BenchmarkMode.TRAIN_AND_EVAL:\n      # Remove "eval" from params so it is not accidentally used. Since eval can\n      # still occur despite params.eval being False, params.eval should never\n      # be used. We cannot yet remove this unconditionally, because the SSD\n      # model still uses params.eval, and hence does not work properly with\n      # --eval_during_training_*.\n      # TODO(b/116627045): We should also remove fields that have an eval\n      # equivalent, like num_batches and num_eval_batches.\n      self.params = remove_param_fields(self.params, {\'eval\'})\n\n  @contextlib.contextmanager\n  def _do_eval(self):\n    """Context manager to switches BenchmarkCNN to eval mode.\n\n    Any evaluation code should be put under this context manager. This context\n    manager switches self._doing_eval to True. It also switches certain\n    attributes, like self.num_batches and self.num_epochs, to be the number of\n    batches and epochs for evaluation respectively\n\n    Yields:\n      Nothing.\n    """\n    # TODO(b/116627045): Find a more general way of switching attributes to the\n    # eval equivalents.\n    old_doing_eval = self._doing_eval\n    old_num_batches = self.num_batches\n    old_num_epochs = self.num_epochs\n    old_batch_size = self.batch_size\n    try:\n      self._doing_eval = True\n      self.num_batches = self.num_eval_batches\n      self.num_epochs = self.num_eval_epochs\n      self.batch_size = self.eval_batch_size\n      self.model.set_batch_size(self.eval_batch_size // self.num_gpus)\n      yield\n    finally:\n      self._doing_eval = old_doing_eval\n      self.num_batches = old_num_batches\n      self.num_epochs = old_num_epochs\n      self.batch_size = old_batch_size\n      self.model.set_batch_size(old_batch_size // self.num_gpus)\n\n  def _config_benchmark_logger(self):\n    """Config the model garden benchmark logger."""\n    model_benchmark_logger = None\n    if self.params.benchmark_log_dir is not None:\n      try:\n        from official.r1.utils.logs import logger as models_logger  # pylint: disable=g-import-not-at-top\n      except ImportError:\n        tf.logging.fatal(\'Please include tensorflow/models to the PYTHONPATH \'\n                         \'in order to use BenchmarkLogger. Configured \'\n                         \'benchmark_log_dir: %s\'\n                         % self.params.benchmark_log_dir)\n        raise\n      model_benchmark_logger = models_logger.BenchmarkFileLogger(\n          self.params.benchmark_log_dir)\n    self.benchmark_logger = model_benchmark_logger\n\n  # TODO(laigd): this changes the global device list which is used everywhere,\n  # consider refactoring it.\n  def reset_devices_for_task(self, task_num, is_local=False):\n    """Used to imitate another task when building a distributed graph."""\n    worker_prefix = (\'/job:localhost\' if is_local else\n                     \'/job:worker/replica:0/task:%s\' % task_num)\n    self.cpu_device = \'%s/cpu:0\' % worker_prefix\n    self.raw_devices = [\n        \'%s/%s:%i\' % (worker_prefix, self.params.device, i)\n        for i in xrange(self.num_gpus)\n    ]\n    self.devices = self.variable_mgr.get_devices()\n\n  def raw_devices_across_tasks(self, is_local=False):\n    """Returns list of raw device names across all tasks."""\n    if is_local:\n      assert self.num_workers == 1\n      return self.raw_devices\n    else:\n      return [\n          \'job:worker/replica:0/task%s/%s:%i\' % (t, self.params.device, i)\n          for t in xrange(self.num_workers)\n          for i in xrange(self.num_gpus)\n      ]\n\n  def print_info(self):\n    """Print basic information."""\n    benchmark_info = self._get_params_info()\n    log_fn(\'Model:       %s\' % self.model.get_model_name())\n    log_fn(\'Dataset:     %s\' % benchmark_info[\'dataset_name\'])\n    log_fn(\'Mode:        %s\' % self.mode)\n    log_fn(\'SingleSess:  %s\' % benchmark_info[\'single_session\'])\n    log_fn(\'Batch size:  %s global\' % (self.batch_size * self.num_workers))\n    log_fn(\'             %s per device\' % (self.batch_size //\n                                           len(self.raw_devices)))\n    if self.batch_group_size > 1:\n      log_fn(\'             %d batches per prepocessing group\' %\n             self.batch_group_size)\n    log_fn(\'Num batches: %d\' % self.num_batches)\n    log_fn(\'Num epochs:  %.2f\' % self.num_epochs)\n    log_fn(\'Devices:     %s\' % benchmark_info[\'device_list\'])\n    log_fn(\'NUMA bind:   %s\' % self.params.use_numa_affinity)\n    log_fn(\'Data format: %s\' % self.params.data_format)\n    if self.rewriter_config:\n      log_fn(\'RewriterConfig: %s\' % self.rewriter_config)\n    log_fn(\'Optimizer:   %s\' % self.params.optimizer)\n    log_fn(\'Variables:   %s\' % self.params.variable_update)\n    if (self.params.variable_update == \'replicated\' or\n        self.params.variable_update == \'distributed_all_reduce\'\n        or self.params.variable_update == \'collective_all_reduce\'):\n      log_fn(\'AllReduce:   %s\' % self.params.all_reduce_spec)\n    if self.job_name:\n      log_fn(\'Sync:        %s\' % self.params.cross_replica_sync)\n    if self.params.staged_vars:\n      log_fn(\'Staged vars: %s\' % self.params.staged_vars)\n    if self.params.variable_update == \'horovod\' and self.params.horovod_device:\n      log_fn(\'Horovod on:  %s\' % self.params.horovod_device)\n    log_fn(\'==========\')\n\n  def _get_params_info(self):\n    """Get the common parameters info for the benchmark run.\n\n    Returns:\n      A dict of processed parameters.\n    """\n    dataset_name = self.dataset.name\n    if self.dataset.use_synthetic_gpu_inputs():\n      dataset_name += \' (synthetic)\'\n    single_session = self.params.variable_update == \'distributed_all_reduce\'\n    if single_session:\n      device_list = self.raw_devices_across_tasks()\n    elif self.params.variable_update == \'horovod\':\n      device_list = [\'horovod/%s:%d\' % (self.params.device, idx)\n                     for idx in range(self.num_workers)]\n    else:\n      device_list = self.raw_devices\n    return {\n        \'dataset_name\': dataset_name,\n        \'single_session\': single_session,\n        \'device_list\': device_list,}\n\n  def _log_benchmark_run(self):\n    """Log the benchmark info to the logger.\n\n    The info logged here should be similar to print_info(), but in a structured\n    JSON format.\n    """\n    if self.benchmark_logger:\n      benchmark_info = self._get_params_info()\n\n      run_param = {\n          \'model\': self.model.get_model_name(),\n          \'dataset\': benchmark_info[\'dataset_name\'],\n          \'mode\': self.mode,\n          \'single_sess\': benchmark_info[\'single_session\'],\n          \'devices\': benchmark_info[\'device_list\'],\n          \'batch_size\': self.batch_size,\n          \'batch_size_per_device\': self.batch_size // len(self.raw_devices),\n          \'num_batches\': self.num_batches,\n          \'num_epochs\': self.num_epochs,\n          \'data_format\': self.params.data_format,\n          \'rewrite_config\': self.rewriter_config,\n          \'optimizer\': self.params.optimizer,\n          \'session_config\': create_config_proto(self.params),\n      }\n      # TODO(scottzhu): tf_cnn_benchmark might execute several times with\n      # different param setting on the same box. This will cause the run file to\n      # only contain the latest info. The benchmark_log_dir should be updated\n      # for every new run.\n      self.benchmark_logger.log_run_info(\n          self.model.get_model_name(), benchmark_info[\'dataset_name\'],\n          run_param, test_id=self.params.benchmark_test_id)\n\n  def run(self):\n    """Run the benchmark task assigned to this process.\n\n    Returns:\n      Dictionary of statistics for training or eval.\n    Raises:\n       ValueError: unrecognized job name.\n    """\n    if self.params.job_name == \'ps\':\n      log_fn(\'Running parameter server %s\' % self.task_index)\n      self.cluster_manager.join_server()\n      return {}\n\n    # For distributed_all_reduce with multiple workers, drive\n    # from a separate controller process.\n    if self.params.variable_update == \'distributed_all_reduce\':\n      if self.params.job_name == \'worker\':\n        log_fn(\'Starting worker %s\' % self.task_index)\n        self.cluster_manager.join_server()\n        return\n      elif self.params.job_name and self.params.job_name != \'controller\':\n        raise ValueError(\'unrecognized job name: %s\' % self.params.job_name)\n\n    self._log_benchmark_run()\n    if self._doing_eval:\n      with tf.Graph().as_default():\n        # TODO(laigd): freeze the graph in eval mode.\n        return self._run_eval()\n    else:\n      return self._benchmark_train()\n\n  def _run_eval(self):\n    """Evaluate a model every self.params.eval_interval_secs.\n\n    Returns:\n      Dictionary containing eval statistics. Currently returns an empty\n      dictionary.\n\n    Raises:\n      ValueError: If self.params.train_dir is unspecified.\n    """\n    if self.params.train_dir is None:\n      raise ValueError(\'Trained model directory not specified\')\n    graph_info = self._build_eval_graph()\n    saver = tf.train.Saver(self.variable_mgr.savable_variables())\n    summary_writer = tf.summary.FileWriter(self.params.eval_dir,\n                                           tf.get_default_graph())\n    target = \'\'\n    # TODO(huangyp): Check if checkpoints haven\'t updated for hours and abort.\n    while True:\n      with tf.Session(\n          target=target, config=create_config_proto(self.params)) as sess:\n        image_producer = None\n        try:\n          global_step = load_checkpoint(saver, sess, self.params.train_dir)\n          image_producer = self._initialize_eval_graph(\n              graph_info.enqueue_ops, graph_info.input_producer_op,\n              graph_info.local_var_init_op_group, sess)\n        except CheckpointNotFoundException:\n          log_fn(\'Checkpoint not found in %s\' % self.params.train_dir)\n        else:  # Only executes if an exception was not thrown\n          self._eval_once(sess, summary_writer, graph_info.fetches,\n                          graph_info.summary_op, image_producer, global_step)\n        if image_producer is not None:\n          image_producer.done()\n        if self.params.eval_interval_secs <= 0:\n          break\n        time.sleep(self.params.eval_interval_secs)\n    return {}\n\n  def _build_eval_graph(self, scope_name=None):\n    """Build the evaluation graph.\n\n    Args:\n      scope_name: String to filter what summaries are collected. Only summary\n        ops whose name contains `scope_name` will be added, which is useful for\n        only including evaluation ops.\n\n    Returns:\n      A GraphInfo named_tuple containing various useful ops and tensors of the\n      evaluation grpah.\n    """\n    with self._do_eval():\n      input_producer_op, enqueue_ops, fetches = self._build_model()\n      local_var_init_op = tf.local_variables_initializer()\n      table_init_ops = tf.tables_initializer()\n      variable_mgr_init_ops = [local_var_init_op]\n      if table_init_ops:\n        variable_mgr_init_ops.extend([table_init_ops])\n      with tf.control_dependencies([local_var_init_op]):\n        variable_mgr_init_ops.extend(self.variable_mgr.get_post_init_ops())\n      local_var_init_op_group = tf.group(*variable_mgr_init_ops)\n\n      summary_op = tf.summary.merge_all(scope=scope_name)\n      # The eval graph has no execution barrier because it doesn\'t run in\n      # distributed mode.\n      execution_barrier = None\n      # We do not use the global step during evaluation.\n      global_step = None\n      return GraphInfo(input_producer_op, enqueue_ops, fetches,\n                       execution_barrier, global_step, local_var_init_op_group,\n                       summary_op)\n\n  # TODO(reedwm): For consistency, we should have a similar\n  # "_initialize_train_graph" function. They can likely be the same function.\n  def _initialize_eval_graph(self, enqueue_ops, input_producer_op,\n                             local_var_init_op_group, sess):\n    """Initializes the evaluation graph.\n\n    Args:\n      enqueue_ops: Ops that adds the preprocessed images to the staging areas.\n      input_producer_op: Op that produce the input batches (before\n        preprocessing).\n      local_var_init_op_group: Group of ops that perform per-device\n        initialization work.\n      sess: The session to initialize the eval graph with.\n\n    Returns:\n      An ImageProducer, or None if an ImageProducer isn\'t being used.\n    """\n    with self._do_eval():\n      if local_var_init_op_group is not None:\n        # We might reinitialize local variables if they were already initialized\n        # during training. This is OK.\n        sess.run(local_var_init_op_group)\n      if self.dataset.queue_runner_required():\n        tf.train.start_queue_runners(sess=sess)\n      image_producer = None\n      if input_producer_op is not None:\n        image_producer = cnn_util.ImageProducer(\n            sess, input_producer_op, self.batch_group_size,\n            self.params.use_python32_barrier)\n        image_producer.start()\n      if enqueue_ops:\n        for i in xrange(len(enqueue_ops)):\n          sess.run(enqueue_ops[:(i + 1)])\n          if image_producer is not None:\n            image_producer.notify_image_consumption()\n      return image_producer\n\n  def _eval_once(self, sess, summary_writer, fetches, summary_op,\n                 image_producer, global_step):\n    """Evaluate the model using the validation dataset."""\n    with self._do_eval():\n      mlperf.logger.log_eval_epoch(\n          mlperf.tags.EVAL_START, global_step, self.batch_size)\n      loop_start_time = start_time = time.perf_counter()\n      # TODO(laigd): refactor the part to compute/report the accuracy. Currently\n      # it only works for image models.\n      top_1_accuracy_sum = 0.0\n      top_5_accuracy_sum = 0.0\n      total_eval_count = self.num_batches * self.batch_size\n      for step in xrange(self.num_batches):\n        if (summary_writer and self.params.save_summaries_steps > 0 and\n            (step + 1) % self.params.save_summaries_steps == 0):\n          results, summary_str = sess.run([fetches, summary_op])\n          summary_writer.add_summary(summary_str)\n        else:\n          results = sess.run(fetches)\n        # Make global_step available in results for postprocessing.\n        results[\'global_step\'] = global_step\n        results = self.model.postprocess(results)\n        top_1_accuracy_sum += results[\'top_1_accuracy\']\n        top_5_accuracy_sum += results[\'top_5_accuracy\']\n        if (step + 1) % self.params.display_every == 0:\n          duration = time.perf_counter() - start_time\n          examples_per_sec = (\n              self.batch_size * self.params.display_every / duration)\n          log_fn(\'%i\\t%.1f examples/sec\' % (step + 1, examples_per_sec))\n          start_time = time.perf_counter()\n        if image_producer is not None:\n          image_producer.notify_image_consumption()\n      loop_end_time = time.perf_counter()\n      accuracy_at_1 = top_1_accuracy_sum / self.num_batches\n      accuracy_at_5 = top_5_accuracy_sum / self.num_batches\n      summary = tf.Summary()\n      summary.value.add(tag=\'eval/Accuracy@1\', simple_value=accuracy_at_1)\n      summary.value.add(tag=\'eval/Accuracy@5\', simple_value=accuracy_at_5)\n      for result_key, result_value in results.items():\n        if result_key.startswith(constants.SIMPLE_VALUE_RESULT_PREFIX):\n          prefix_len = len(constants.SIMPLE_VALUE_RESULT_PREFIX)\n          summary.value.add(tag=\'eval/\' + result_key[prefix_len:],\n                            simple_value=result_value)\n      if summary_writer:\n        summary_writer.add_summary(summary, global_step)\n      log_fn(\'Accuracy @ 1 = %.4f Accuracy @ 5 = %.4f [%d examples]\' %\n             (accuracy_at_1, accuracy_at_5, total_eval_count))\n      elapsed_time = loop_end_time - loop_start_time\n      images_per_sec = (self.num_batches * self.batch_size / elapsed_time)\n      if self.mode != constants.BenchmarkMode.TRAIN_AND_EVAL:\n        # Note that we compute the top 1 accuracy and top 5 accuracy for each\n        # batch, which will have a slight performance impact.\n        log_fn(\'-\' * 64)\n        log_fn(\'total images/sec: %.2f\' % images_per_sec)\n        log_fn(\'-\' * 64)\n      if self.benchmark_logger:\n        eval_result = {\n            \'eval_top_1_accuracy\', accuracy_at_1,\n            \'eval_top_5_accuracy\', accuracy_at_5,\n            \'eval_average_examples_per_sec\', images_per_sec,\n            tf.GraphKeys.GLOBAL_STEP, global_step,\n        }\n        self.benchmark_logger.log_evaluation_result(eval_result)\n      mlperf.logger.log_eval_epoch(\n          mlperf.tags.EVAL_STOP, global_step, self.batch_size)\n      mlperf.logger.log(key=mlperf.tags.EVAL_SIZE,\n                        value=self.num_batches * self.batch_size)\n      if self.params.model != \'ssd300\':  # ssd300 logs eval accuracy elsewhere.\n        mlperf.logger.log_eval_accuracy(\n            accuracy_at_1, global_step, self.train_batch_size,\n            examples_per_epoch=self.dataset.num_examples_per_epoch(\'train\'))\n      if self.params.stop_at_top_1_accuracy:\n        mlperf.logger.log(key=mlperf.tags.EVAL_TARGET,\n                          value=self.params.stop_at_top_1_accuracy)\n      return accuracy_at_1, accuracy_at_5\n\n  def _benchmark_train(self):\n    """Run cnn in benchmark mode. Skip the backward pass if forward_only is on.\n\n    Returns:\n      Dictionary containing training statistics (num_workers, num_steps,\n      average_wall_time, images_per_sec).\n    """\n    graph = tf.Graph()\n    with graph.as_default():\n      build_result = self._build_graph()\n      if self.mode == constants.BenchmarkMode.TRAIN_AND_EVAL:\n        with self.variable_mgr.reuse_variables():\n          with tf.name_scope(\'Evaluation\') as ns:\n            eval_build_results = self._build_eval_graph(ns)\n      else:\n        eval_build_results = None\n    (graph, result_to_benchmark) = self._preprocess_graph(graph, build_result)\n    with graph.as_default():\n      return self._benchmark_graph(result_to_benchmark, eval_build_results)\n\n  GPU_CACHED_INPUT_VARIABLE_NAME = \'gpu_cached_inputs\'\n\n  def _unfreezable_local_variables(self, graph):\n    """Get the local variables that we don\'t want to freeze."""\n    return graph.get_collection(\n        tf.GraphKeys.LOCAL_VARIABLES,\n        # We don\'t freeze the gpu_cached_images local variable so it won\'t get\n        # constant folded with ops which process the input.\n        scope=\'.*\' + BenchmarkCNN.GPU_CACHED_INPUT_VARIABLE_NAME)\n\n  def _build_graph(self):\n    """Build the graph.\n\n    Returns:\n      A namedtuple containing the ops/tensors that required by\n      _benchmark_graph().\n    """\n    if self.single_session:\n      (input_producer_op, enqueue_ops, fetches) = (\n          self._build_model_single_session())\n    else:\n      (input_producer_op, enqueue_ops, fetches) = self._build_model()\n    fetches_list = nest.flatten(list(fetches.values()))\n    main_fetch_group = tf.group(*fetches_list, name=\'main_fetch_group\')\n    execution_barrier = None\n    if (not self.single_session and self.job_name and\n        not self.params.cross_replica_sync):\n      execution_barrier = self.add_sync_queues_and_barrier(\n          \'execution_barrier_\', [])\n\n    global_step = tf.train.get_global_step()\n    with tf.device(self.global_step_device), tf.name_scope(\'inc_global_step\'):\n      with tf.control_dependencies([main_fetch_group]):\n        fetches[\'inc_global_step\'] = global_step.assign_add(1)\n        images = tf.reshape(\n            images,\n            shape=[\n                self.batch_size, image_size, image_size,\n                self.dataset.depth\n            ])\n    else:\n    # Skips the init ops for freezable local variables in forward_only mode so\n    # we can remove all the assign ops when converting variables to constants.\n    with tf.name_scope(\'local_variable_initialization\'):\n      if self.forward_only_and_freeze:\n        local_var_init_op = tf.variables_initializer(\n            self._unfreezable_local_variables(tf.get_default_graph()))\n      else:\n        local_var_init_op = tf.local_variables_initializer()\n    table_init_ops = tf.tables_initializer()\n\n    variable_manager_init_ops = [local_var_init_op]\n    if table_init_ops:\n      variable_manager_init_ops.extend([table_init_ops])\n    if not self.forward_only_and_freeze:\n      with tf.control_dependencies([local_var_init_op]):\n        variable_manager_init_ops.extend(self.variable_mgr.get_post_init_ops())\n    if ((not self.single_session) and (not self.distributed_collective) and\n        self.job_name and self.params.cross_replica_sync):\n      # Ensure all workers execute variable_manager_init_ops before they start\n      # executing the model.\n      variable_manager_init_ops.append(\n          self.add_sync_queues_and_barrier(\'init_ops_end_\',\n                                           variable_manager_init_ops))\n    local_var_init_op_group = tf.group(*variable_manager_init_ops,\n                                       name=\'local_var_init_op_group\')\n    summary_op = tf.summary.merge_all()\n\n    return GraphInfo(\n        input_producer_op=input_producer_op,\n        enqueue_ops=enqueue_ops,\n        fetches=fetches,\n        execution_barrier=execution_barrier,\n        global_step=global_step,\n        local_var_init_op_group=local_var_init_op_group,\n        summary_op=summary_op)\n\n  def _benchmark_graph(self, graph_info, eval_graph_info):\n    """Benchmark the training graph.\n\n    Args:\n      graph_info: the namedtuple returned by _build_graph() which\n        contains all necessary information to benchmark the graph, including\n        named tensors/ops list, fetches, etc.\n      eval_graph_info: Similar to graph_info but for the eval graph if\n        --eval_during_training_* is used. Otherwise, None.\n    Returns:\n      Dictionary containing training statistics (num_workers, num_steps,\n      average_wall_time, images_per_sec).\n    """\n    log_fn(\'Initializing graph\')\n    if self.params.variable_update == \'horovod\':\n      import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n      # First worker will be \'chief\' - it will write summaries and\n      # save checkpoints.\n      is_chief = hvd.rank() == 0\n    else:\n      is_chief = (not self.job_name or self.task_index == 0)\n\n    summary_writer = None\n    if (is_chief and self.params.summary_verbosity and self.params.train_dir and\n        self.params.save_summaries_steps > 0):\n      summary_writer = tf.summary.FileWriter(self.params.train_dir,\n                                             tf.get_default_graph())\n\n    # We want to start the benchmark timer right after a image_producer barrier\n    # and avoids undesired waiting times on barriers.\n    if ((self.num_warmup_batches + len(graph_info.enqueue_ops) - 1) %\n        self.batch_group_size) != 0:\n      self.num_warmup_batches = int(\n          math.ceil(\n              (self.num_warmup_batches + len(graph_info.enqueue_ops) - 1.0) /\n              (self.batch_group_size)) * self.batch_group_size -\n          len(graph_info.enqueue_ops) + 1)\n      log_fn(\'Round up warm up steps to %d to match batch_group_size\' %\n             self.num_warmup_batches)\n      assert ((self.num_warmup_batches + len(graph_info.enqueue_ops) - 1) %\n              self.batch_group_size) == 0\n    # We run the summaries in the same thread as the training operations by\n    # passing in None for summary_op to avoid a summary_thread being started.\n    # Running summaries and training operations in parallel could run out of\n    # GPU memory.\n    if is_chief and not self.forward_only_and_freeze:\n      saver = tf.train.Saver(\n          self.variable_mgr.savable_variables(),\n          save_relative_paths=True,\n          max_to_keep=self.params.max_ckpts_to_keep)\n    else:\n      saver = None\n    ready_for_local_init_op = None\n    if self.job_name and not (self.single_session or\n                              self.distributed_collective):\n      # In distributed mode, we don\'t want to run local_var_init_op_group until\n      # the global variables are initialized, because local_var_init_op_group\n      # may use global variables (such as in distributed replicated mode). We\n      # don\'t set this in non-distributed mode, because in non-distributed mode,\n      # local_var_init_op_group may itself initialize global variables (such as\n      # in replicated mode).\n      ready_for_local_init_op = tf.report_uninitialized_variables(\n          tf.global_variables())\n    if self.params.variable_update == \'horovod\':\n      import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n      bcast_global_variables_op = hvd.broadcast_global_variables(0)\n    else:\n      bcast_global_variables_op = None\n\n    if self.params.variable_update == \'collective_all_reduce\':\n      # It doesn\'t matter what this collective_graph_key value is,\n      # so long as it\'s > 0 and the same at every worker.\n      init_run_options = tf.RunOptions()\n      init_run_options.experimental.collective_graph_key = 6\n    else:\n      init_run_options = tf.RunOptions()\n    local_var_init_ops = [graph_info.local_var_init_op_group]\n    if eval_graph_info:\n      # `eval_graph_info.local_var_init_op_group` also includes some of the\n      # training initializer ops, since it\'s difficult to filter them out.\n      # Rerunning the training initializer ops is OK, but we add a control\n      # dependency since running two sets of training initializer ops at the\n      # same time can cause race conditions.\n      with tf.control_dependencies(local_var_init_ops):\n        local_var_init_ops.append(eval_graph_info.local_var_init_op_group)\n    sv = tf.train.Supervisor(\n        # For the purpose of Supervisor, all Horovod workers are \'chiefs\',\n        # since we want session to be initialized symmetrically on all the\n        # workers.\n        is_chief=is_chief or (self.params.variable_update == \'horovod\'\n                              or self.distributed_collective),\n        # Log dir should be unset on non-chief workers to prevent Horovod\n        # workers from corrupting each other\'s checkpoints.\n        logdir=self.params.train_dir if is_chief else None,\n        ready_for_local_init_op=ready_for_local_init_op,\n        local_init_op=local_var_init_ops,\n        saver=saver,\n        global_step=graph_info.global_step,\n        summary_op=None,\n        save_model_secs=self.params.save_model_secs,\n        summary_writer=summary_writer,\n        local_init_run_options=init_run_options)\n\n    profiler = tf.profiler.Profiler() if self.params.tfprof_file else None\n    if self.graph_file is not None:\n      path, filename = os.path.split(self.graph_file)\n      as_text = filename.endswith(\'txt\')\n      log_fn(\'Writing GraphDef as %s to %s\' % (  # pyformat break\n          \'text\' if as_text else \'binary\', self.graph_file))\n      tf.train.write_graph(tf.get_default_graph().as_graph_def(add_shapes=True),\n                           path, filename, as_text)\n\n    start_standard_services = (\n        self.params.train_dir or\n        self.dataset.queue_runner_required())\n    target = self.cluster_manager.get_target() if self.cluster_manager else \'\'\n    with sv.managed_session(\n        master=target,\n        config=create_config_proto(self.params),\n        start_standard_services=start_standard_services) as sess:\n      # Anything that can potentially raise an OutOfRangeError with \'sess\' MUST\n      # be under this try block. The managed_session() context manager silently\n      # ignores OutOfRangeError, so we must catch them and wrap them with\n      # a different exception type so that they can be propagated up to the\n      # caller.\n      try:\n        stats = self.benchmark_with_session(\n            sess, sv, graph_info, eval_graph_info, bcast_global_variables_op,\n            is_chief, summary_writer, profiler)\n      except tf.errors.OutOfRangeError:\n        raise RuntimeError(\n            \'Received OutOfRangeError. Wrapping in Runtime error to avoid \'\n            \'Supervisor from suppressing the error. Original OutOfRangeError \'\n            \'with traceback:\\n\' + traceback.format_exc())\n\n    sv.stop()\n    if profiler:\n      generate_tfprof_profile(profiler, self.params.tfprof_file)\n    return stats\n\n  def benchmark_with_session(self, sess, supervisor, graph_info,\n                             eval_graph_info, bcast_global_variables_op,\n                             is_chief, summary_writer, profiler):\n    """Benchmarks the graph with the given session.\n\n    Args:\n      sess: The session to benchmark the graph with\n      supervisor: The Supervisor that created the session.\n      graph_info: the namedtuple returned by _build_graph() which\n        contains all necessary information to benchmark the graph, including\n        named tensors/ops list, fetches, etc.\n      eval_graph_info: Similar to graph_info but for the eval graph if\n        --eval_during_training_every_n_steps is used. Otherwise, None.\n      bcast_global_variables_op: If Horovod is used, the op to broadcast the\n        global variables to all the processes. None if Horovod is not used.\n      is_chief: True if this is the chief process.\n      summary_writer: The SummaryWriter used to write summaries, or None if\n        summaries are not used.\n      profiler: The tf.profiler.Profiler, or None if tfprof is not used.\n\n    Returns:\n      Dictionary containing training statistics (num_workers, num_steps,\n      average_wall_time, images_per_sec).\n    """\n    if self.params.backbone_model_path is not None:\n      self.model.load_backbone_model(sess, self.params.backbone_model_path)\n    if bcast_global_variables_op:\n      sess.run(bcast_global_variables_op)\n    image_producer = None\n    if graph_info.input_producer_op is not None:\n      image_producer = cnn_util.ImageProducer(\n          sess, graph_info.input_producer_op, self.batch_group_size,\n          self.params.use_python32_barrier)\n      image_producer.start()\n    if graph_info.enqueue_ops:\n      for i in xrange(len(graph_info.enqueue_ops)):\n        sess.run(graph_info.enqueue_ops[:(i + 1)])\n        if image_producer is not None:\n          image_producer.notify_image_consumption()\n    self.init_global_step, = sess.run([graph_info.global_step])\n    if self.job_name and not self.params.cross_replica_sync:\n      # TODO(zhengxq): Do we need to use a global step watcher at all?\n      global_step_watcher = GlobalStepWatcher(\n          sess, graph_info.global_step,\n          self.num_workers * self.num_warmup_batches +\n          self.init_global_step,\n          self.num_workers * (self.num_warmup_batches + self.num_batches) - 1)\n      global_step_watcher.start()\n    else:\n      global_step_watcher = None\n    eval_image_producer = None\n    if eval_graph_info:\n      # We pass local_var_init_op_group=None because the Supervisor already\n      # initialized local variables above. We need to have the Supervisor\n      # initialize the local variables, because otherwise it throws an error\n      # complaining that not all variables were initialized.\n      eval_image_producer = self._initialize_eval_graph(\n          eval_graph_info.enqueue_ops, eval_graph_info.input_producer_op,\n          local_var_init_op_group=None, sess=sess)\n    step_train_times = []\n    log_fn(\'Running warm up\')\n    local_step = -1 * self.num_warmup_batches\n    if self.single_session:\n      # In single session mode, each step, the global_step is incremented by\n      # 1. In non-single session mode, each step, the global_step is\n      # incremented once per worker. This means we need to divide\n      # init_global_step by num_workers only in non-single session mode.\n      end_local_step = self.num_batches - self.init_global_step\n    else:\n      end_local_step = self.num_batches - (self.init_global_step //\n                                           self.num_workers)\n    if not global_step_watcher:\n      # In cross-replica sync mode, all workers must run the same number of\n      # local steps, or else the workers running the extra step will block.\n      done_fn = lambda: local_step >= end_local_step\n    else:\n      done_fn = global_step_watcher.done\n    if self.params.debugger is not None:\n      if self.params.debugger == \'cli\':\n        log_fn(\'The CLI TensorFlow debugger will be used.\')\n        sess = tf_debug.LocalCLIDebugWrapperSession(sess)\n      else:\n        log_fn(\'The TensorBoard debugger plugin will be used.\')\n        sess = tf_debug.TensorBoardDebugWrapperSession(sess,\n                                                       self.params.debugger)\n    mlperf.logger.log(key=mlperf.tags.TRAIN_LOOP)\n    skip_final_eval = False\n    accuracy_at_1 = None\n    accuracy_at_5 = None\n    last_eval_step = local_step\n    loop_start_time = time.perf_counter()\n    last_average_loss = None\n    while not done_fn():\n      if local_step == 0:\n        log_fn(\'Done warm up\')\n        if graph_info.execution_barrier:\n          log_fn(\'Waiting for other replicas to finish warm up\')\n          sess.run([graph_info.execution_barrier])\n\n        # TODO(laigd): rename \'Img\' to maybe \'Input\'.\n        header_str = (\'Step\\tImg/sec\\t\' +\n                      self.params.loss_type_to_report.replace(\'/\', \' \'))\n        if self.params.print_training_accuracy or self.params.forward_only:\n          # TODO(laigd): use the actual accuracy op names of the model.\n          header_str += \'\\ttop_1_accuracy\\ttop_5_accuracy\'\n        log_fn(header_str)\n        assert len(step_train_times) == self.num_warmup_batches\n        # reset times to ignore warm up batch\n        step_train_times = []\n        loop_start_time = time.perf_counter()\n      if (summary_writer and\n          (local_step + 1) % self.params.save_summaries_steps == 0):\n        fetch_summary = graph_info.summary_op\n      else:\n        fetch_summary = None\n      collective_graph_key = 7 if (\n          self.params.variable_update == \'collective_all_reduce\') else 0\n      (summary_str, last_average_loss) = benchmark_one_step(\n          sess, graph_info.fetches, local_step,\n          self.batch_size * (self.num_workers\n                             if self.single_session else 1), step_train_times,\n          self.trace_filename, self.params.partitioned_graph_file_prefix,\n          profiler, image_producer, self.params, fetch_summary,\n          benchmark_logger=self.benchmark_logger,\n          collective_graph_key=collective_graph_key,\n          should_output_files=(self.params.variable_update != \'horovod\' or\n                               is_chief))\n      if summary_str is not None and is_chief:\n        supervisor.summary_computed(sess, summary_str)\n      local_step += 1\n      if (self.params.save_model_steps and\n          local_step % self.params.save_model_steps == 0 and\n          local_step > 0 and\n          is_chief):\n        supervisor.saver.save(sess, supervisor.save_path,\n                              supervisor.global_step)\n      if (eval_graph_info and local_step > 0 and not done_fn() and\n          self._should_eval_during_training(local_step)):\n        python_global_step = sess.run(graph_info.global_step)\n        num_steps_since_last_eval = local_step - last_eval_step\n        # The INPUT_SIZE tag value might not match the\n        # PREPROC_NUM_TRAIN_EXAMPLES tag value, because the number of examples\n        # run, which is INPUT_SIZE, is rounded up to the nearest multiple of\n        # self.batch_size.\n        mlperf.logger.log(\n            key=mlperf.tags.INPUT_SIZE,\n            value=num_steps_since_last_eval * self.batch_size)\n        log_fn(\'Running evaluation at global_step {}\'.format(\n            python_global_step))\n        accuracy_at_1, accuracy_at_5 = self._eval_once(\n            sess, summary_writer, eval_graph_info.fetches,\n            eval_graph_info.summary_op, eval_image_producer,\n            python_global_step)\n        last_eval_step = local_step\n        if (self.params.stop_at_top_1_accuracy and\n            accuracy_at_1 >= self.params.stop_at_top_1_accuracy):\n          log_fn(\'Stopping, as eval accuracy at least %s was reached\' %\n                 self.params.stop_at_top_1_accuracy)\n          skip_final_eval = True\n          break\n        else:\n          log_fn(\'Resuming training\')\n      if eval_graph_info and self.model.reached_target():\n        log_fn(\'Stopping, as the model indicates its custom goal was reached\')\n        skip_final_eval = True\n        break\n    loop_end_time = time.perf_counter()\n    # Waits for the global step to be done, regardless of done_fn.\n    if global_step_watcher:\n      while not global_step_watcher.done():\n        time.sleep(.25)\n    if not global_step_watcher:\n      elapsed_time = loop_end_time - loop_start_time\n      average_wall_time = elapsed_time / local_step if local_step > 0 else 0\n      images_per_sec = (self.num_workers * local_step * self.batch_size /\n                        elapsed_time)\n      num_steps = local_step * self.num_workers\n    else:\n      # NOTE: Each worker independently increases the global step. So,\n      # num_steps will be the sum of the local_steps from each worker.\n      num_steps = global_step_watcher.num_steps()\n      elapsed_time = global_step_watcher.elapsed_time()\n      average_wall_time = (elapsed_time * self.num_workers / num_steps\n                           if num_steps > 0 else 0)\n      images_per_sec = num_steps * self.batch_size / elapsed_time\n\n    # We skip printing images/sec if --eval_during_training_* is specified,\n    # because we are both processing training and evaluation images, so a\n    # singular "images/sec" value is meaningless.\n    if self.mode != constants.BenchmarkMode.TRAIN_AND_EVAL:\n      log_fn(\'-\' * 64)\n      # TODO(laigd): rename \'images\' to maybe \'inputs\'.\n      log_fn(\'total images/sec: %.2f\' % images_per_sec)\n      log_fn(\'-\' * 64)\n    else:\n      log_fn(\'Done with training\')\n    num_steps_since_last_eval = local_step - last_eval_step\n    mlperf.logger.log(\n        key=mlperf.tags.INPUT_SIZE,\n        value=num_steps_since_last_eval * self.batch_size)\n    python_global_step = sess.run(graph_info.global_step)\n    if eval_graph_info and not skip_final_eval:\n      log_fn(\'Running final evaluation at global_step {}\'.format(\n          python_global_step))\n      accuracy_at_1, accuracy_at_5 = self._eval_once(\n          sess, summary_writer, eval_graph_info.fetches,\n          eval_graph_info.summary_op, eval_image_producer, python_global_step)\n    num_epochs_ran = (python_global_step * self.batch_size /\n                      self.dataset.num_examples_per_epoch(\'train\'))\n    mlperf.logger.log_train_epochs(num_epochs_ran)\n    if image_producer is not None:\n      image_producer.done()\n    if eval_image_producer is not None:\n      eval_image_producer.done()\n    if is_chief:\n      if self.benchmark_logger:\n        self.benchmark_logger.log_metric(\n            \'average_examples_per_sec\', images_per_sec, global_step=num_steps)\n\n    # Save the model checkpoint.\n    if self.params.train_dir is not None and is_chief:\n      checkpoint_path = os.path.join(self.params.train_dir, \'model.ckpt\')\n      if not gfile.Exists(self.params.train_dir):\n        gfile.MakeDirs(self.params.train_dir)\n      supervisor.saver.save(sess, checkpoint_path, graph_info.global_step)\n    if graph_info.execution_barrier:\n      # Wait for other workers to reach the end, so this worker doesn\'t\n      # go away underneath them.\n      sess.run([graph_info.execution_barrier])\n    stats = {\n        \'num_workers\': self.num_workers,\n        \'num_steps\': num_steps,\n        \'average_wall_time\': average_wall_time,\n        \'images_per_sec\': images_per_sec\n    }\n    if last_average_loss is not None:\n      stats[\'last_average_loss\'] = last_average_loss\n    if accuracy_at_1 is not None:\n      stats[\'top_1_accuracy\'] = accuracy_at_1\n    if accuracy_at_5 is not None:\n      stats[\'top_5_accuracy\'] = accuracy_at_5\n\n    success = bool(self.model.reached_target() or\n                   (accuracy_at_1 and self.params.stop_at_top_1_accuracy and\n                    accuracy_at_1 >= self.params.stop_at_top_1_accuracy))\n    mlperf.logger.log(key=mlperf.tags.RUN_STOP, value={\'success\': success})\n    mlperf.logger.log(key=mlperf.tags.RUN_FINAL)\n    return stats\n\n  def _should_eval_during_training(self, step):\n    """Return True iff should run eval during training at current step."""\n\n    assert self.mode == constants.BenchmarkMode.TRAIN_AND_EVAL\n\n    if self.params.eval_during_training_every_n_steps:\n      return step % self.params.eval_during_training_every_n_steps == 0\n\n    # All other --eval_during_training_* flags are converted to step numbers\n    # at which the model should run evaluation during training.\n    return step in self.eval_during_training_at_specified_steps\n\n  def _preprocess_graph(self, graph, graph_info):\n    """Preprocess the graph before executing.\n\n    Depending on the params, it runs various preprocessing on the graph,\n    including freezing, TensorRT conversion, etc.\n\n    Args:\n      graph: the graph to preprocess.\n      graph_info: the namedtuple returned by _build_graph() which\n        contains all necessary information to benchmark the graph, including\n        named tensors/ops list, fetches, etc.\n\n    Returns:\n      The updated graph and graph_info with the ops/tensors/fetches updated\n      according to the imported graph.\n    """\n    assert isinstance(graph_info.fetches, dict)\n    assert isinstance(graph_info.global_step, tf.Variable)\n    if not self.forward_only_and_freeze:\n      return (graph, graph_info)\n\n    # Get the names of the ops that need to keep during conversion.\n    flattened_op_names = list(\n        set([\n            v.name.split(\':\')[0]\n            for v in nest.flatten(graph_info)\n            if v is not None\n        ]))\n    # Get variables that we don\'t want to freeze.\n    # Only keep unfreezable variables in forward_only_and_freeze mode.\n    # TODO(laigd): consider making global_step a constant.\n    variables_to_keep = {graph_info.global_step: tf.GraphKeys.GLOBAL_VARIABLES}\n    variables_to_keep.update({\n        local_variable: tf.GraphKeys.LOCAL_VARIABLES\n        for local_variable in self._unfreezable_local_variables(graph)\n    })\n\n    variable_initializers = [\n        variable.initializer.name for variable in variables_to_keep]\n    output_node_names = (\n        flattened_op_names +\n        # Add variable initializer and read ops to the output list, so\n        # convert_variables_to_constants() will keep them.\n        variable_initializers +\n        [variable.value().op.name for variable in variables_to_keep])\n    graphdef = graph.as_graph_def(add_shapes=True)\n\n    # Freeze the graph.\n    with graph.as_default():\n      with tf.Session(config=create_config_proto(self.params)) as sess:\n        sess.run(tf.global_variables_initializer())\n        sess.run(tf.local_variables_initializer())\n        graphdef = graph_util.convert_variables_to_constants(\n            sess,\n            graphdef,\n            output_node_names,\n            variable_names_blacklist=[\n                variable.op.name for variable in variables_to_keep\n            ])\n\n    # Run TensorRT conversion.\n    if self.params.trt_mode:\n      # Import here instead of at top, because this will crash if TensorRT is\n      # not installed\n      from tensorflow.python.compiler.tensorrt import trt_convert  # pylint: disable=g-import-not-at-top\n      # Avoid TF-TRT bridge from touching all variable initializer ops and their\n      # dependencies, since they can directly be fetched by sess.run()s that\n      # initialize the variables.\n      # pylint: disable=protected-access\n      name_to_input_name, _, _ = graph_util_impl._extract_graph_summary(\n          graphdef)\n      initializer_subgraph_ops = graph_util_impl._bfs_for_reachable_nodes(\n          variable_initializers, name_to_input_name)\n      # pylint: enable=protected-access\n\n      graphdef = trt_convert.create_inference_graph(\n          graphdef,\n          outputs=output_node_names + list(initializer_subgraph_ops),\n          max_batch_size=self.model.get_batch_size(),\n          max_workspace_size_bytes=self.params.trt_max_workspace_size_bytes,\n          precision_mode=self.params.trt_mode)\n\n    # Creates a new graph as the default and import the converted graph back.\n    updated_graph = tf.Graph()\n\n    def _get_tensors_or_ops(inputs):\n      """Gets the updated tensors or ops from \'updated_graph\'."""\n\n      def _get_fn(element):\n        if element is None:\n          return None\n        if \':\' in element.name:\n          return updated_graph.get_tensor_by_name(element.name)\n        return updated_graph.get_operation_by_name(element.name)\n\n      if isinstance(inputs, (list, dict, tuple)):\n        return nest.map_structure(_get_fn, inputs)\n      else:\n        return _get_fn(inputs)\n\n    with updated_graph.as_default():\n      importer.import_graph_def(graph_def=graphdef, name=\'\')\n\n      # Update the variables\n      for variable in variables_to_keep:\n        updated_variable = tf.Variable.from_proto(variable.to_proto())\n        tf.add_to_collection(variables_to_keep[variable], updated_variable)\n        if variable is graph_info.global_step:\n          updated_global_step = updated_variable\n\n    updated_graph_info = GraphInfo(\n        input_producer_op=_get_tensors_or_ops(graph_info.input_producer_op),\n        enqueue_ops=_get_tensors_or_ops(graph_info.enqueue_ops),\n        execution_barrier=_get_tensors_or_ops(graph_info.execution_barrier),\n        local_var_init_op_group=_get_tensors_or_ops(\n            graph_info.local_var_init_op_group),\n        fetches=_get_tensors_or_ops(graph_info.fetches),\n        global_step=updated_global_step,\n        summary_op=None)\n    return (updated_graph, updated_graph_info)\n\n  def _build_input_processing(self, shift_ratio=0):\n    """"Build the image (pre)processing portion of the model graph.\n\n    Args:\n      shift_ratio: shift_ratio for data_flow_ops.RecordInput.\n\n    Returns:\n      An InputProcessingInfo containing all the input sources to the model.\n    """\n    input_processing_info = InputProcessingInfo(\n        input_producer_op=None,\n        input_producer_stages=None,\n        multi_device_iterator_input=None)\n\n    mlperf.logger.log(key=mlperf.tags.INPUT_ORDER)\n    if not self._doing_eval:\n      mlperf.logger.log(key=mlperf.tags.INPUT_BATCH_SIZE, value=self.batch_size)\n\n    # If using synthetic gpu inputs, do nothing on the cpu side.\n    if self.dataset.use_synthetic_gpu_inputs():\n      assert not self.datasets_use_prefetch\n      return input_processing_info\n\n    if self._doing_eval:\n      input_preprocessor = self.eval_input_preprocessor\n      mlperf.logger.log(key=mlperf.tags.PREPROC_NUM_EVAL_EXAMPLES,\n                        value=self.dataset.num_examples_per_epoch(\'validation\'))\n    else:\n      input_preprocessor = self.input_preprocessor\n      mlperf.logger.log(key=mlperf.tags.PREPROC_NUM_TRAIN_EXAMPLES,\n                        value=self.dataset.num_examples_per_epoch(\'train\'))\n\n    # Use prefetching mechanism provided by dataset input pipeline.\n    if self.datasets_use_prefetch:\n      multi_device_iterator = (\n          input_preprocessor.build_multi_device_iterator(\n              self.batch_size, len(self.devices), self.cpu_device, self.params,\n              self.raw_devices, self.dataset, self._doing_eval))\n      return input_processing_info._replace(\n          multi_device_iterator_input=multi_device_iterator.get_next())\n\n    # Not using dataset prefetching. Use a staging area to mimic the prefetching\n    # behavior instead.\n    with tf.device(self.cpu_device):\n      if self._doing_eval:\n        subset = \'validation\'\n      else:\n        subset = \'train\'\n      input_list = input_preprocessor.minibatch(\n          self.dataset,\n          subset=subset,\n          params=self.params,\n          shift_ratio=shift_ratio)\n\n      input_producer_op = []\n      input_producer_stages = []\n      for device_num in range(len(self.devices)):\n        staging_area = data_flow_ops.StagingArea(\n            [parts[0].dtype for parts in input_list],\n            shapes=[parts[0].get_shape() for parts in input_list],\n            shared_name=\'input_producer_staging_area_%d_eval_%s\' %\n            (device_num, self._doing_eval))\n        input_producer_stages.append(staging_area)\n        for group_index in xrange(self.batch_group_size):\n          batch_index = group_index + device_num * self.batch_group_size\n          put_op = staging_area.put(\n              [parts[batch_index] for parts in input_list])\n          input_producer_op.append(put_op)\n      assert input_producer_op\n\n    return input_processing_info._replace(\n        input_producer_op=input_producer_op,\n        input_producer_stages=input_producer_stages)\n\n  def _maybe_initialize_fp16(self):\n    """Initialize fp16 settings."""\n    if self.params.use_fp16 and not self._doing_eval:\n      init_loss_scale_val = float(self.params.fp16_loss_scale or\n                                  self.model.get_fp16_loss_scale())\n      self.loss_scale = None\n      self.loss_scale_normal_steps = None\n      if self.enable_auto_loss_scale or init_loss_scale_val != 1:\n        self.loss_scale = tf.get_variable(\n            name=\'loss_scale\',\n            initializer=init_loss_scale_val,\n            dtype=tf.float32,\n            trainable=False)\n      if self.enable_auto_loss_scale:\n        self.loss_scale_normal_steps = tf.get_variable(\n            name=\'loss_scale_normal_steps\', initializer=0, trainable=False)\n\n  def _build_model(self):\n    """Build the TensorFlow graph."""\n    if self.datasets_use_prefetch:\n      assert not self.params.staged_vars\n      assert not self.variable_mgr.supports_staged_vars()\n\n    # Adjust seed so different workers start read different input files.\n    if self.params.variable_update == \'horovod\':\n      import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n      seed_adjustment = hvd.rank()\n    else:\n      seed_adjustment = 0\n    mlperf.logger.log(key=mlperf.tags.RUN_SET_RANDOM_SEED,\n                      value=self.params.tf_random_seed + seed_adjustment)\n    tf.set_random_seed(self.params.tf_random_seed + seed_adjustment)\n    mlperf.logger.log(key=mlperf.tags.RUN_SET_RANDOM_SEED,\n                      value=4321 + seed_adjustment)\n    np.random.seed(4321 + seed_adjustment)\n    phase_train = not (self._doing_eval or self.params.forward_only)\n\n    if self._doing_eval:\n      mode_string = \'evaluation\'\n    else:\n      mode_string = \'training\'\n\n    log_fn(\'Generating {} model\'.format(mode_string))\n    losses = []\n    device_grads = []\n    all_logits = []\n    all_accuracy_ops = {}\n    gpu_compute_stage_ops = []\n    gpu_grad_stage_ops = []\n\n    with tf.device(self.global_step_device):\n      global_step = tf.train.get_or_create_global_step()\n      self._maybe_initialize_fp16()\n\n    # Build the processing and model for the worker.\n    input_producer_op = None\n    with tf.name_scope(\'input_processing\'):\n      input_processing_info = self._build_input_processing(shift_ratio=0)\n      if input_processing_info.input_producer_op is not None:\n        input_producer_op = tf.group(*input_processing_info.input_producer_op)\n    update_ops = None\n    staging_delta_ops = []\n\n    for device_num in range(len(self.devices)):\n      with tf.name_scope(\'tower_%i\' % device_num) as name_scope, (\n          self.variable_mgr.create_outer_variable_scope(device_num)):\n        results = self.add_forward_pass_and_gradients(\n            phase_train, device_num, device_num, input_processing_info,\n            gpu_compute_stage_ops, gpu_grad_stage_ops)\n\n        if self.params.backbone_model_path:\n          self.model.add_backbone_saver()\n\n        if phase_train:\n          losses.append(results[\'loss\'])\n          device_grads.append(results[\'gradvars\'])\n        else:\n          all_logits.append(results[\'logits\'])\n        if not phase_train or self.params.print_training_accuracy:\n          for name, op in results.items():\n            if name.startswith(\'accuracy:\'):\n              key = name[9:]\n              if key not in all_accuracy_ops:\n                all_accuracy_ops[key] = []\n              all_accuracy_ops[key].append(op)\n\n        if device_num == 0:\n          # Retain the Batch Normalization updates operations only from the\n          # first tower. These operations update the moving mean and moving\n          # variance variables, which are updated (but not used) during\n          # training, and used during evaluation. The moving mean and variance\n          # approximate the true mean and variance across all images in the\n          # dataset. Therefore, in replicated mode, these moving averages would\n          # be almost identical for each tower, and so we only update and save\n          # the moving averages for one tower. In parameter server mode, all\n          # towers share a copy of the variables so we also only need to update\n          # and save the moving averages once.\n          update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS, name_scope)\n          if self.datasets_use_prefetch:\n            assert not self.variable_mgr.staging_delta_ops\n          else:\n            staging_delta_ops = list(self.variable_mgr.staging_delta_ops)\n\n    enqueue_ops = []\n    if not self.datasets_use_prefetch:\n      if self.variable_mgr.supports_staged_vars():\n        for staging_ops in self.variable_mgr.staging_vars_on_devices:\n          gpu_compute_stage_ops.extend(\n              [put_op for _, (put_op, _) in six.iteritems(staging_ops)])\n      enqueue_ops.append(tf.group(*gpu_compute_stage_ops,\n                                  name=\'gpu_compute_stage_ops_group\'))\n      if gpu_grad_stage_ops:\n        staging_delta_ops += gpu_grad_stage_ops\n      if staging_delta_ops:\n        enqueue_ops.append(tf.group(*(staging_delta_ops)))\n\n    if (self.mode == constants.BenchmarkMode.TRAIN_AND_EVAL and\n        self.params.variable_update == \'replicated\'):\n      # We need to get all the update ops instead of only those for the first\n      # tower. This is because during evaluation, each tower will read from its\n      # own tower\'s moving averages instead of the first tower\'s moving\n      # averages.\n      # TODO(reedwm): Have each tower read from the first tower\'s moving\n      # averages for a slight performance gain.\n      update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n      mlperf.logger.log(key=mlperf.tags.INPUT_BN_SPAN,\n                        value=self.batch_size // len(self.raw_devices))\n\n    fetches = self._build_fetches(global_step, all_logits, losses, device_grads,\n                                  enqueue_ops, update_ops, all_accuracy_ops,\n                                  phase_train)\n    return (input_producer_op, enqueue_ops, fetches)\n\n  def _build_fetches(self, global_step, all_logits, losses, device_grads,\n                     enqueue_ops, update_ops, all_accuracy_ops, phase_train):\n    """Complete construction of model graph, populating the fetches map."""\n    fetches = {}\n    if enqueue_ops:\n      fetches[\'enqueue_ops\'] = enqueue_ops\n    for name, ops in all_accuracy_ops.items():\n      # For fetches that starts with \'tensor:\', keep dimension and skip reducing\n      # them to scalars.\n      if name.startswith(constants.UNREDUCED_ACCURACY_OP_PREFIX):\n        key = name[len(constants.UNREDUCED_ACCURACY_OP_PREFIX):]\n        fetches[key] = tf.concat(ops, 0)\n      else:\n        fetches[name] = (\n            tf.reduce_sum(ops) /\n            (self.batch_size *\n             (self.num_workers if self.single_session else 1)))\n        if self.task_index == 0 and self.params.summary_verbosity >= 1:\n          tf.summary.scalar(name, fetches[name])\n\n    if not phase_train:\n      if self.params.forward_only:\n        fetches[\'all_logits\'] = tf.concat(all_logits, 0)\n      return fetches\n    apply_gradient_devices, gradient_state = (\n        self.variable_mgr.preprocess_device_grads(device_grads))\n\n    # TODO(reedwm): Greatly simplify the learning rate code.\n    if (self.params.variable_update == \'horovod\' or\n        self.params.variable_update == \'collective_all_reduce\'):\n      # Each worker independently increments global_step.\n      examples_per_step = self.batch_size * self.num_workers\n    else:\n      # global_step is shared by all workers, and so every iteration\n      # global_step is incremented by num_workers.\n      examples_per_step = self.batch_size\n    if self.params.compute_lr_on_cpu:\n      with tf.device(self.cpu_device):\n        learning_rate = get_learning_rate(self.params, global_step,\n                                          self.dataset.num_examples_per_epoch(),\n                                          self.model, examples_per_step)\n\n    training_ops = []\n    for d, device in enumerate(apply_gradient_devices):\n      with tf.device(device):\n        with tf.name_scope(\'average_loss\'):\n          average_loss = tf.reduce_mean(losses)\n        with tf.name_scope(\'get_gradients_to_apply\'):\n          avg_grads = self.variable_mgr.get_gradients_to_apply(d,\n                                                               gradient_state)\n\n        if not self.params.compute_lr_on_cpu:\n          # We compute the learning rate once for each device in\n          # `apply_gradient_devices`.\n          learning_rate = get_learning_rate(\n              self.params, global_step, self.dataset.num_examples_per_epoch(),\n              self.model, examples_per_step)\n        gradient_clip = self.params.gradient_clip\n        if gradient_clip is not None:\n          with tf.name_scope(\'clip_gradients\'):\n            clipped_grads = [(tf.clip_by_value(grad, -gradient_clip,\n                                               +gradient_clip), var)\n                             for grad, var in avg_grads]\n        else:\n          clipped_grads = avg_grads\n\n        learning_rate = tf.identity(learning_rate, name=\'learning_rate_tensor\')\n        opt = get_optimizer(self.params, learning_rate)\n        loss_scale_params = variable_mgr_util.AutoLossScaleParams(\n            enable_auto_loss_scale=self.enable_auto_loss_scale,\n            loss_scale=self.loss_scale,\n            loss_scale_normal_steps=self.loss_scale_normal_steps,\n            inc_loss_scale_every_n=self.params.fp16_inc_loss_scale_every_n,\n            is_chief=not self.job_name or self.task_index == 0)\n\n        with tf.name_scope(\'append_apply_gradient_ops\'):\n          self.variable_mgr.append_apply_gradients_ops(\n              gradient_state, opt, clipped_grads, training_ops,\n              loss_scale_params)\n    train_op = tf.group(*(training_ops + update_ops), name=\'train_ops_group\')\n\n    with tf.device(self.cpu_device):\n      if self.task_index == 0 and self.params.summary_verbosity >= 1:\n        tf.summary.scalar(\'learning_rate\', learning_rate)\n        tf.summary.scalar(self.params.loss_type_to_report, average_loss)\n        if self.loss_scale is not None:\n          tf.summary.scalar(\'loss_scale\', self.loss_scale)\n        if self.loss_scale_normal_steps:\n          tf.summary.scalar(\'loss_scale_normal_steps\',\n                            self.loss_scale_normal_steps)\n\n        if self.params.summary_verbosity >= 2:\n          self.gradient_histogram_summary(avg_grads)\n\n        if self.params.summary_verbosity >= 3:\n          for grad, var in avg_grads:\n            if grad is not None:\n              tf.summary.histogram(var.op.name + \'/gradients\', grad)\n          for var in tf.trainable_variables():\n            tf.summary.histogram(var.op.name, var)\n\n    fetches[\'train_op\'] = train_op\n    fetches[\'average_loss\'] = average_loss\n    return fetches\n\n  def gradient_histogram_summary(self, avg_grads):\n    """Create histogram of log values of all non-zero gradients."""\n    with tf.name_scope(\'log_gradients_summary\'):\n      all_grads = []\n      for grad, _ in avg_grads:\n        all_grads.append(tf.reshape(grad, [-1]))\n      grads = tf.abs(tf.concat(all_grads, 0))\n      # exclude grads with zero values.\n      indices_for_non_zero_grads = tf.where(tf.not_equal(grads, 0))\n      log_grads = tf.reshape(\n          tf.log(tf.gather(grads, indices_for_non_zero_grads)), [-1])\n      tf.summary.histogram(\'log_gradients\', log_grads)\n\n  def _build_model_single_session(self):\n    """Build the TensorFlow graph for multiple replicas in a single_session.\n\n    Returns:\n      input_producer_op:\n      enqueue_ops:\n      fetches:\n\n    Raises:\n       ValueError: optimizer not recognized.\n\n    Single session runs multiple model replicas as part of one large\n    distributed graph, whose global execution is always step-synchronized.\n    """\n    # verify assumptions\n    assert self.params.task_index == 0\n    assert not self._doing_eval\n    assert not self.params.forward_only\n    assert not self.params.staged_vars\n\n    tf.set_random_seed(self.params.tf_random_seed)\n    np.random.seed(4321)\n    phase_train = True\n\n    log_fn(\'Generating training model\')\n    losses = []\n    device_grads = []\n    all_logits = []\n    all_accuracy_ops = {}\n    gpu_compute_stage_ops = []\n    gpu_grad_stage_ops = []\n\n    with tf.device(self.global_step_device):\n      global_step = tf.train.get_or_create_global_step()\n\n    update_ops = []\n    global_input_producer_op = []\n\n    is_local = not self.job_name\n    if is_local:\n      assert self.num_workers == 1\n    for task_num in range(self.num_workers):\n      # Reset the devices that self.variable_mgr knows about to those\n      # belonging to the next worker (task).\n      self.reset_devices_for_task(task_num, is_local)\n      # Build the per-worker image processing\n      with tf.name_scope(\'input_processing\'):\n        input_processing_info = self._build_input_processing(\n            shift_ratio=(task_num / self.num_workers))\n      if input_processing_info.input_producer_op is not None:\n        global_input_producer_op.extend(input_processing_info.input_producer_op)\n      # Build the per-worker model replica.\n      for rel_device_num in range(len(self.devices)):\n        abs_device_num = task_num * len(self.devices) + rel_device_num\n        with self.variable_mgr.create_outer_variable_scope(\n            abs_device_num), tf.name_scope(\n                \'task_%i_tower_%i\' % (task_num, rel_device_num)) as name_scope:\n          task_results = self.add_forward_pass_and_gradients(\n              phase_train, rel_device_num, abs_device_num,\n              input_processing_info, gpu_compute_stage_ops, gpu_grad_stage_ops)\n\n          if self.params.backbone_model_path:\n            self.model.add_backbone_saver()\n\n          if phase_train:\n            losses.append(task_results[\'loss\'])\n            device_grads.append(task_results[\'gradvars\'])\n          else:\n            all_logits.append(task_results[\'logits\'])\n          if not phase_train or self.params.print_training_accuracy:\n            for name, op in task_results.items():\n              if name.startswith(\'accuracy:\'):\n                key = name[9:]\n                if key not in all_accuracy_ops:\n                  all_accuracy_ops[key] = []\n                all_accuracy_ops[key].append(op)\n\n          if rel_device_num == 0:\n            # Retain the Batch Normalization updates operations only\n            # from the first tower. These operations update the moving\n            # mean and moving variance variables, which are updated\n            # (but not used) during training, and used during\n            # evaluation. The moving mean and variance approximate the\n            # true mean and variance across all images in the\n            # dataset. Therefore, in replicated mode, these moving\n            # averages would be almost identical for each tower, and\n            # so we only update and save the moving averages for one\n            # tower. In parameter server mode, all towers share a copy\n            # of the variables so we also only need to update and save\n            # the moving averages once.\n            update_ops.extend(\n                tf.get_collection(tf.GraphKeys.UPDATE_OPS, name_scope))\n            assert not self.variable_mgr.staging_delta_ops\n\n    enqueue_ops = []\n    if gpu_compute_stage_ops:\n      enqueue_ops.append(tf.group(*gpu_compute_stage_ops,\n                                  name=\'gpu_compute_stage_ops\'))\n    assert not self.variable_mgr.supports_staged_vars()\n    assert not gpu_grad_stage_ops\n\n    fetches = self._build_fetches(global_step, all_logits, losses, device_grads,\n                                  enqueue_ops, update_ops, all_accuracy_ops,\n                                  phase_train)\n    if global_input_producer_op:\n      global_input_producer_op = tf.group(*global_input_producer_op)\n    else:\n      global_input_producer_op = None\n    return (global_input_producer_op, enqueue_ops, fetches)\n\n  def add_forward_pass_and_gradients(self,\n                                     phase_train,\n                                     rel_device_num,\n                                     abs_device_num,\n                                     input_processing_info,\n                                     gpu_compute_stage_ops,\n                                     gpu_grad_stage_ops):\n    """Add ops for forward-pass and gradient computations."""\n    nclass = self.dataset.num_classes\n    if self.datasets_use_prefetch:\n      assert input_processing_info.multi_device_iterator_input, (\n          \'multi_device_iterator_input cannot be None if \'\n          \'datasets_use_prefetch=True\')\n      input_list = (\n          input_processing_info.multi_device_iterator_input[rel_device_num])\n    else:\n      if not self.dataset.use_synthetic_gpu_inputs():\n        input_producer_stage = input_processing_info.input_producer_stages[\n            rel_device_num]\n        with tf.device(self.cpu_device):\n          host_input_list = input_producer_stage.get()\n        with tf.device(self.raw_devices[rel_device_num]):\n          gpu_compute_stage = data_flow_ops.StagingArea(\n              [inp.dtype for inp in host_input_list],\n              shapes=[inp.get_shape() for inp in host_input_list])\n          # The CPU-to-GPU copy is triggered here.\n          gpu_compute_stage_op = gpu_compute_stage.put(host_input_list)\n          input_list = gpu_compute_stage.get()\n          gpu_compute_stage_ops.append(gpu_compute_stage_op)\n      else:\n        with tf.device(self.raw_devices[rel_device_num]):\n          # Minor hack to avoid H2D copy when using synthetic data\n          input_list = self.model.get_synthetic_inputs(\n              BenchmarkCNN.GPU_CACHED_INPUT_VARIABLE_NAME, nclass)\n\n    # Labels reshaping happens all on gpu:0. Reshaping synthetic labels on\n    # multiple devices slows down XLA computation for an unknown reason.\n    # TODO(b/116875203): Find/address root cause of XLA slow down.\n    labels_device_placement_hack = (\n        self.dataset.use_synthetic_gpu_inputs() and self.params.xla_compile)\n\n    def device_aware_reshape(tensor, shape):\n      device = self.devices[rel_device_num]\n      # Labels are int32, place reshapes on gpu:0 (no device placement) when the\n      # hack is enabled.\n      if labels_device_placement_hack and tensor.dtype == tf.int32:\n        device = \'\'\n      with tf.device(device):\n        return tf.reshape(tensor, shape=shape)\n\n    subset = \'validation\' if self._doing_eval else \'train\'\n    input_shapes = self.model.get_input_shapes(subset)\n    input_list = [\n        device_aware_reshape(input_list[i], shape=input_shapes[i])\n        for i in range(len(input_list))\n    ]\n\n    def forward_pass_and_gradients():\n      """Builds forward pass and gradient computation network.\n\n      When phase_train=True and print_training_accuracy=False:\n        return [loss] + grads\n\n      When phase_train=True and print_training_accuracy=True:\n        return [logits, loss] + grads\n\n      When phase_train=False,\n        return [logits]\n\n      Its output can always be unpacked by\n\n      ```\n        outputs = forward_pass_and_gradients()\n        logits, loss, grads = unpack_forward_pass_and_gradients_output(outputs)\n      ```\n\n      Returns:\n        outputs: A list of tensors depending on different modes.\n      """\n\n      build_network_result = self.model.build_network(\n          input_list, phase_train, nclass)\n      logits = build_network_result.logits\n\n      if not phase_train:\n        return [logits]\n\n      base_loss = self.model.loss_function(input_list, build_network_result)\n      params = self.variable_mgr.trainable_variables_on_device(\n          rel_device_num, abs_device_num)\n      l2_loss = None\n      total_loss = base_loss\n      with tf.name_scope(\'l2_loss\'):\n        fp32_params = params\n        if self.model.data_type == tf.float16 and self.params.fp16_vars:\n          # fp16 reductions are very slow on GPUs, so cast to fp32 before\n          # calling tf.nn.l2_loss and tf.add_n.\n          # TODO(b/36217816): Once the bug is fixed, investigate if we should do\n          # this reduction in fp16.\n          fp32_params = (tf.cast(p, tf.float32) for p in params)\n        filtered_params = self.model.filter_l2_loss_vars(fp32_params)\n        if rel_device_num == len(self.devices) - 1:\n          # We compute the L2 loss for only one device instead of all of them,\n          # because the L2 loss for each device is the same. To adjust for this,\n          # we multiply the L2 loss by the number of devices. We choose the\n          # last device because for some reason, on a Volta DGX1, the first four\n          # GPUs take slightly longer to complete a step than the last four.\n          # TODO(reedwm): Shard the L2 loss computations across GPUs.\n          if self.params.single_l2_loss_op:\n            # TODO(reedwm): If faster, create a fused op that does the L2 loss\n            # on multiple tensors, and use that instead of concatenating\n            # tensors.\n            reshaped_params = [tf.reshape(p, (-1,)) for p in filtered_params]\n            l2_loss = tf.nn.l2_loss(tf.concat(reshaped_params, axis=0))\n          else:\n            l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in filtered_params])\n      weight_decay = self.params.weight_decay\n      mlperf.logger.log(key=mlperf.tags.OPT_WEIGHT_DECAY, value=weight_decay)\n      if (weight_decay is not None and weight_decay != 0. and\n          l2_loss is not None):\n        mlperf.logger.log(key=mlperf.tags.MODEL_L2_REGULARIZATION,\n                          value=weight_decay)\n        total_loss += len(self.devices) * weight_decay * l2_loss\n\n      aggmeth = tf.AggregationMethod.DEFAULT\n      scaled_loss = (total_loss if self.loss_scale is None\n                     else total_loss * self.loss_scale)\n      grads = tf.gradients(scaled_loss, params, aggregation_method=aggmeth)\n      if self.params.sparse_to_dense_grads:\n        # Passing a sparse gradient to convert_to_tensor turns it into a dense\n        # gradient. A sparse gradient is an instance of tf.IndexedSlices.\n        # convert_to_tensor does not modify dense tensors.\n        grads = [tf.convert_to_tensor(g) for g in grads]\n      if self.loss_scale is not None:\n        # TODO(reedwm): If automatic loss scaling is not used, we could avoid\n        # these multiplications by directly modifying the learning rate instead.\n        # If this is done, care must be taken to ensure that this scaling method\n        # is correct, as some optimizers square gradients and do other\n        # operations which might not be compatible with modifying both the\n        # gradients and the learning rate.\n\n        grads = [\n            grad * tf.cast(1. / self.loss_scale, grad.dtype) for grad in grads\n        ]\n\n      if self.params.variable_update == \'horovod\':\n        import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n        if self.params.horovod_device:\n          horovod_device = \'/%s:0\' % self.params.horovod_device\n        else:\n          horovod_device = \'\'\n        # All-reduce gradients using Horovod.\n        grads = [hvd.allreduce(grad, average=False, device_dense=horovod_device)\n                 for grad in grads]\n\n      if self.params.staged_vars:\n        grad_dtypes = [grad.dtype for grad in grads]\n        grad_shapes = [grad.shape for grad in grads]\n        grad_stage = data_flow_ops.StagingArea(grad_dtypes, grad_shapes)\n        grad_stage_op = grad_stage.put(grads)\n        # In general, this decouples the computation of the gradients and\n        # the updates of the weights.\n        # During the pipeline warm up, this runs enough training to produce\n        # the first set of gradients.\n        gpu_grad_stage_ops.append(grad_stage_op)\n        grads = grad_stage.get()\n\n      if self.params.loss_type_to_report == \'total_loss\':\n        loss = total_loss\n      else:\n        loss = base_loss\n\n      if self.params.print_training_accuracy:\n        return [logits, loss] + grads\n      else:\n        return [loss] + grads\n\n    def unpack_forward_pass_and_gradients_output(forward_pass_and_grad_outputs):\n      """Unpacks outputs from forward_pass_and_gradients.\n\n      Args:\n        forward_pass_and_grad_outputs: Output from forward_pass_and_gradients.\n\n      Returns:\n        logits: Unscaled probability distribution from forward pass.\n          If unavailable, None is returned.\n        loss: Loss function result from logits.\n          If unavailable, None is returned.\n        grads: Gradients for all trainable variables.\n          If unavailable, None is returned.\n      """\n      logits = None\n      # logits is only fetched in non-train mode or when\n      # print_training_accuracy is set.\n      if not phase_train or self.params.print_training_accuracy:\n        logits = forward_pass_and_grad_outputs.pop(0)\n\n      loss = (\n          forward_pass_and_grad_outputs[0]\n          if forward_pass_and_grad_outputs else None)\n      grads = (\n          forward_pass_and_grad_outputs[1:]\n          if forward_pass_and_grad_outputs else None)\n\n      return logits, loss, grads\n\n    def make_results(logits, loss, grads):\n      """Generate results based on logits, loss and grads."""\n      results = {}  # The return value\n\n      if logits is not None:\n        results[\'logits\'] = logits\n        accuracy_ops = self.model.accuracy_function(input_list, logits)\n        for name, op in accuracy_ops.items():\n          results[\'accuracy:\' + name] = op\n\n      if loss is not None:\n        results[\'loss\'] = loss\n\n      if grads is not None:\n        param_refs = self.variable_mgr.trainable_variables_on_device(\n            rel_device_num, abs_device_num, writable=True)\n        results[\'gradvars\'] = list(zip(grads, param_refs))\n\n      return results\n\n    with tf.device(self.devices[rel_device_num]):\n      outputs = maybe_compile(forward_pass_and_gradients, self.params)\n      logits, loss, grads = unpack_forward_pass_and_gradients_output(outputs)\n      return make_results(logits, loss, grads)\n\n  def get_input_preprocessor(self):\n    """Returns the image preprocessor to used, based on the model.\n\n    Returns:\n      The image preprocessor, or None if synthetic data should be used.\n    """\n    shift_ratio = 0\n    if self.job_name:\n      # shift_ratio prevents multiple workers from processing the same batch\n      # during a step\n      shift_ratio = self.task_index / self.num_workers\n\n    processor_class = self.dataset.get_input_preprocessor(\n        self.params.input_preprocessor)\n    assert processor_class\n    subset = \'validation\' if self._doing_eval else \'train\'\n    return processor_class(\n        self.batch_size * self.batch_group_size,\n        self.model.get_input_shapes(subset),\n        len(self.devices) * self.batch_group_size,\n        dtype=self.model.data_type,\n        train=(not self._doing_eval),\n        # TODO(laigd): refactor away image model specific parameters.\n        distortions=self.params.distortions,\n        resize_method=self.resize_method,\n        shift_ratio=shift_ratio,\n        summary_verbosity=self.params.summary_verbosity,\n        distort_color_in_yiq=self.params.distort_color_in_yiq,\n        fuse_decode_and_crop=self.params.fuse_decode_and_crop,\n        match_mlperf=self.params.ml_perf)\n\n  def add_sync_queues_and_barrier(self, name_prefix, enqueue_after_list):\n    """Adds ops to enqueue on all worker queues.\n\n    Args:\n      name_prefix: prefixed for the shared_name of ops.\n      enqueue_after_list: control dependency from ops.\n\n    Returns:\n      An op that should be used as control dependency before starting next step.\n    """\n    self.sync_queue_counter += 1\n    with tf.device(self.sync_queue_devices[(\n        self.sync_queue_counter % len(self.sync_queue_devices))]):\n      sync_queues = [\n          tf.FIFOQueue(self.num_workers, [tf.bool], shapes=[[]],\n                       shared_name=\'%s%s\' % (name_prefix, i))\n          for i in range(self.num_workers)]\n      queue_ops = []\n      # For each other worker, add an entry in a queue, signaling that it can\n      # finish this step.\n      token = tf.constant(False)\n      with tf.control_dependencies(enqueue_after_list):\n        for i, q in enumerate(sync_queues):\n          if i == self.task_index:\n            queue_ops.append(tf.no_op())\n          else:\n            queue_ops.append(q.enqueue(token))\n\n      # Drain tokens off queue for this worker, one for each other worker.\n      queue_ops.append(\n          sync_queues[self.task_index].dequeue_many(len(sync_queues) - 1))\n\n      return tf.group(*queue_ops)\n\n\ndef _is_mkl_flag_absent(mkl_flag):\n  return not (absl_flags.FLAGS.is_parsed() and mkl_flag in absl_flags.FLAGS\n              and absl_flags.FLAGS[mkl_flag].present)\n\n\ndef _print_os_env_ignored_warning(mkl_flag, flag_default_val, os_env_var):\n  tf.logging.warn(\n      (\'OS ENV variable %s=%s is ignored and script default: \'\n       \'%s is used. Use --%s to override.\') %\n      (os_env_var, os.environ[os_env_var], flag_default_val, mkl_flag))\n\n\ndef set_default_param_values_and_env_vars(params):\n  """Sets up the default param values and environment variables ."""\n  if params.batchnorm_persistent:\n    os.environ[\'TF_USE_CUDNN_BATCHNORM_SPATIAL_PERSISTENT\'] = \'1\'\n  else:\n    os.environ.pop(\'TF_USE_CUDNN_BATCHNORM_SPATIAL_PERSISTENT\', None)\n  if params.winograd_nonfused:\n    os.environ[\'TF_ENABLE_WINOGRAD_NONFUSED\'] = \'1\'\n  else:\n    os.environ.pop(\'TF_ENABLE_WINOGRAD_NONFUSED\', None)\n  if params.autotune_threshold:\n    os.environ[\'TF_AUTOTUNE_THRESHOLD\'] = str(params.autotune_threshold)\n  os.environ[\'TF_SYNC_ON_FINISH\'] = str(int(params.sync_on_finish))\n  argparse.ArgumentParser(\n      formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n\n  # Sets environment variables for MKL\n  # If OS ENV vars are overridden by script defaults, a warning msg is printed.\n  if params.mkl:\n    mkl_flags = [\'kmp_blocktime\', \'kmp_settings\', \'kmp_affinity\',\n                 \'num_intra_threads\']\n    for mkl_flag in mkl_flags:\n      os_env_var = mkl_flag.upper()\n      if mkl_flag == \'num_intra_threads\':\n        os_env_var = \'OMP_NUM_THREADS\'\n      flag_val = str(getattr(params, mkl_flag))\n      if _is_mkl_flag_absent(mkl_flag) and os_env_var in os.environ:\n        _print_os_env_ignored_warning(mkl_flag, flag_val, os_env_var)\n      os.environ[os_env_var] = flag_val\n      if mkl_flag == \'num_intra_threads\' and not params.num_intra_threads:\n        os.environ.pop(os_env_var, None)\n\n  # Sets GPU thread settings\n  if params.device.lower() == \'gpu\':\n    params = params._replace(gpu_thread_mode=params.gpu_thread_mode.lower())\n    if params.gpu_thread_mode not in [\'global\', \'gpu_shared\', \'gpu_private\']:\n      raise ValueError(\'Invalid gpu_thread_mode: %s\' % params.gpu_thread_mode)\n    os.environ[\'TF_GPU_THREAD_MODE\'] = params.gpu_thread_mode\n\n    if params.per_gpu_thread_count and params.gpu_thread_mode == \'global\':\n      raise ValueError(\n          \'Invalid per_gpu_thread_count with gpu_thread_mode=global: %s\' %\n          params.per_gpu_thread_count)\n    # Default to two threads. One for the device compute and the other for\n    # memory copies.\n    per_gpu_thread_count = params.per_gpu_thread_count or 2\n    total_gpu_thread_count = per_gpu_thread_count * params.num_gpus\n\n    if params.gpu_thread_mode == \'gpu_private\':\n      os.environ[\'TF_GPU_THREAD_COUNT\'] = str(per_gpu_thread_count)\n    elif params.gpu_thread_mode == \'gpu_shared\':\n      os.environ[\'TF_GPU_THREAD_COUNT\'] = str(total_gpu_thread_count)\n\n    cpu_count = multiprocessing.cpu_count()\n    if not params.num_inter_threads and params.gpu_thread_mode in [\n        \'gpu_private\', \'gpu_shared\'\n    ]:\n      main_thread_count = max(cpu_count - total_gpu_thread_count, 1)\n      params = params._replace(num_inter_threads=main_thread_count)\n\n    if (params.datasets_use_prefetch and\n        params.datasets_num_private_threads is None):\n      # From the total cpu thread count, subtract the total_gpu_thread_count,\n      # and then 2 threads per GPU device for event monitoring and sending /\n      # receiving tensors\n      num_monitoring_threads = 2 * params.num_gpus\n      num_private_threads = max(\n          cpu_count - total_gpu_thread_count - num_monitoring_threads, 1)\n      params = params._replace(datasets_num_private_threads=num_private_threads)\n  return params\n\n\ndef setup(params):\n  """Sets up the environment that BenchmarkCNN should run in.\n\n  Args:\n    params: Params tuple, typically created by make_params or\n      make_params_from_flags.\n\n  Returns:\n    A potentially modified params.\n  Raises:\n    ValueError: invalid parames combinations.\n  """\n  # Set up environment variables before doing any other global initialization to\n  # make sure it uses the appropriate environment variables.\n  params = set_default_param_values_and_env_vars(params)\n\n  # horovod needs to be initialized before create_config_proto() call since\n  # it will be used in config generation if enabled.\n  if params.variable_update == \'horovod\':\n    import horovod.tensorflow as hvd  # pylint: disable=g-import-not-at-top\n    hvd.init()\n\n  platforms_util.initialize(params, create_config_proto(params))\n\n  if not params.job_name:\n    # Create a dummy session to initialize TF global variables using the input\n    # params. Otherwise, ListDevices function may create global devices using\n    # the default config instead of using the user provided config.\n    #\n    # TODO(hinsu): Find a way to achieve the same for distributed benchmark. It\n    # is not legal to create distributed session after local session. It is also\n    # not possible to create distributed session here as that results in\n    # multiple creation of ClusterManager and Server.\n    with tf.Session(config=create_config_proto(params)) as sess:\n      del sess\n\n  return params\n\n\ndef maybe_compile(computation, params):\n  if params and params.xla_compile:\n    return tf.xla.experimental.compile(computation)\n  else:\n    return computation()\n', 'message': 'Fixing the per-gpu batch size.\n\nPiperOrigin-RevId: 190181787\n', 'author_name': 'Reed', 'lang': 'py'}