# https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py # changes: commented out logging # changes: convert_to_tensor_v2 -> convert_to_tensor # changes: extend from tf.python.keras.optimizer_v2.learning_rate_schedule.LearningRateScheduler # Copyright 2018 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. # ============================================================================== """Enable Layer-wise Adaptive Rate Scaling optimizer in ResNet.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function from absl import flags import tensorflow as tf #from tf2_common.utils.mlp_log import mlp_log from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import math_ops from tensorflow.python.keras.optimizer_v2 import learning_rate_schedule FLAGS = flags.FLAGS def define_lars_flags(): """Defines flags needed by LARS optimizer.""" flags.DEFINE_float( 'end_learning_rate', default=None, help=('Polynomial decay end learning rate.')) flags.DEFINE_float( 'lars_epsilon', default=0.0, help=('Override autoselected LARS epsilon.')) flags.DEFINE_float( 'warmup_epochs', default=None, help=('Override autoselected polynomial decay warmup epochs.')) flags.DEFINE_float( 'momentum', default=0.9, help=('Momentum parameter used in the MomentumOptimizer.')) class PolynomialDecayWithWarmup(learning_rate_schedule.LearningRateSchedule): """A LearningRateSchedule that uses a polynomial decay with warmup.""" def __init__( self, batch_size, steps_per_epoch, train_steps, initial_learning_rate=None, end_learning_rate=None, warmup_epochs=None, compute_lr_on_cpu=False, name=None): """Applies a polynomial decay to the learning rate with warmup.""" super(PolynomialDecayWithWarmup, self).__init__() self.batch_size = batch_size self.steps_per_epoch = steps_per_epoch self.train_steps = train_steps self.name = name self.learning_rate_ops_cache = {} self.compute_lr_on_cpu = compute_lr_on_cpu if batch_size < 16384: self.initial_learning_rate = 10.0 warmup_epochs_ = 5 elif batch_size < 32768: self.initial_learning_rate = 25.0 warmup_epochs_ = 5 else: self.initial_learning_rate = 31.2 warmup_epochs_ = 25 # Override default poly learning rate and warmup epochs if initial_learning_rate: self.initial_learning_rate = initial_learning_rate if end_learning_rate: self.end_learning_rate = end_learning_rate else: self.end_learning_rate = 0.0001 if warmup_epochs is not None: warmup_epochs_ = warmup_epochs self.warmup_epochs = warmup_epochs_ """ opt_name = FLAGS.optimizer.lower() mlp_log.mlperf_print('opt_name', opt_name) if opt_name == 'lars': mlp_log.mlperf_print('{}_epsilon'.format(opt_name), FLAGS.lars_epsilon) mlp_log.mlperf_print('{}_opt_weight_decay'.format(opt_name), FLAGS.weight_decay) mlp_log.mlperf_print('{}_opt_base_learning_rate'.format(opt_name), self.initial_learning_rate) mlp_log.mlperf_print('{}_opt_learning_rate_warmup_epochs'.format(opt_name), warmup_epochs_) mlp_log.mlperf_print('{}_opt_end_learning_rate'.format(opt_name), self.end_learning_rate) """ warmup_steps = warmup_epochs_ * steps_per_epoch self.warmup_steps = tf.cast(warmup_steps, tf.float32) self.decay_steps = train_steps - warmup_steps + 1 """ mlp_log.mlperf_print('{}_opt_learning_rate_decay_steps'.format(opt_name), int(self.decay_steps)) mlp_log.mlperf_print( '{}_opt_learning_rate_decay_poly_power'.format(opt_name), 2.0) mlp_log.mlperf_print('{}_opt_momentum'.format(opt_name), FLAGS.momentum) """ self.poly_rate_scheduler = tf.keras.optimizers.schedules.PolynomialDecay( initial_learning_rate=self.initial_learning_rate, decay_steps=self.decay_steps, end_learning_rate=self.end_learning_rate, power=2.0) def __call__(self, step): if tf.executing_eagerly(): return self._get_learning_rate(step) # In an eager function or graph, the current implementation of optimizer # repeatedly call and thus create ops for the learning rate schedule. To # avoid this, we cache the ops if not executing eagerly. graph = tf.compat.v1.get_default_graph() if graph not in self.learning_rate_ops_cache: if self.compute_lr_on_cpu: with tf.device('/device:CPU:0'): self.learning_rate_ops_cache[graph] = self._get_learning_rate(step) else: self.learning_rate_ops_cache[graph] = self._get_learning_rate(step) return self.learning_rate_ops_cache[graph] def _get_learning_rate(self, step): with ops.name_scope_v2(self.name or 'PolynomialDecayWithWarmup') as name: initial_learning_rate = ops.convert_to_tensor( self.initial_learning_rate, name='initial_learning_rate') warmup_steps = ops.convert_to_tensor( self.warmup_steps, name='warmup_steps') warmup_rate = ( initial_learning_rate * step / warmup_steps) poly_steps = math_ops.subtract(step, warmup_steps) poly_rate = self.poly_rate_scheduler(poly_steps) decay_rate = tf.where(step <= warmup_steps, warmup_rate, poly_rate, name=name) return decay_rate def get_config(self): return { 'batch_size': self.batch_size, 'steps_per_epoch': self.steps_per_epoch, 'train_steps': self.train_steps, 'initial_learning_rate': self.initial_learning_rate, 'end_learning_rate': self.end_learning_rate, 'warmup_epochs': self.warmup_epochs, 'name': self.name, }