def build_optimizer( self, lr: Union[tf.keras.optimizers.schedules.LearningRateSchedule, float]): """Build optimizer. Builds optimizer from config. It takes learning rate as input, and builds the optimizer according to the optimizer config. Typically, the learning rate built using self.build_lr() is passed as an argument to this method. Args: lr: A floating point value, or a tf.keras.optimizers.schedules.LearningRateSchedule instance. Returns: tf.keras.optimizers.Optimizer instance. """ optimizer_dict = self._optimizer_config.as_dict() ## Delete clipnorm and clipvalue if None if optimizer_dict['clipnorm'] is None: del optimizer_dict['clipnorm'] if optimizer_dict['clipvalue'] is None: del optimizer_dict['clipvalue'] optimizer_dict['learning_rate'] = lr optimizer = OPTIMIZERS_CLS[self._optimizer_type](**optimizer_dict) if self._use_ema: optimizer = ema_optimizer.ExponentialMovingAverage( optimizer, **self._ema_config.as_dict()) return optimizer
def build_optimizer( self, lr: Union[tf.keras.optimizers.schedules.LearningRateSchedule, float], gradient_transformers: Optional[List[ Callable[[List[Tuple[tf.Tensor, tf.Tensor]]], List[Tuple[tf.Tensor, tf.Tensor]]]]] = None, postprocessor: Optional[ Callable[[tf.keras.optimizers.Optimizer], tf.keras.optimizers.Optimizer]] = None): """Build optimizer. Builds optimizer from config. It takes learning rate as input, and builds the optimizer according to the optimizer config. Typically, the learning rate built using self.build_lr() is passed as an argument to this method. Args: lr: A floating point value, or a tf.keras.optimizers.schedules.LearningRateSchedule instance. gradient_transformers: Optional list of functions to use to transform gradients before applying updates to Variables. The functions are applied after gradient_aggregator. The functions should accept and return a list of (gradient, variable) tuples. clipvalue, clipnorm, global_clipnorm should not be set when gradient_transformers is passed. postprocessor: An optional function for postprocessing the optimizer. It takes an optimizer and returns an optimizer. Returns: `tf.keras.optimizers.Optimizer` or `tf.keras.optimizers.experimental.Optimizer` instance. """ optimizer_dict = self._optimizer_config.as_dict() ## Delete clipnorm, clipvalue, global_clipnorm if None if optimizer_dict['clipnorm'] is None: del optimizer_dict['clipnorm'] if optimizer_dict['clipvalue'] is None: del optimizer_dict['clipvalue'] if optimizer_dict['global_clipnorm'] is None: del optimizer_dict['global_clipnorm'] optimizer_dict['learning_rate'] = lr if gradient_transformers is not None: optimizer_dict['gradient_transformers'] = gradient_transformers optimizer = OPTIMIZERS_CLS[self._optimizer_type](**optimizer_dict) if self._use_ema: optimizer = ema_optimizer.ExponentialMovingAverage( optimizer, **self._ema_config.as_dict()) if postprocessor: optimizer = postprocessor(optimizer) assert isinstance( optimizer, (tf.keras.optimizers.Optimizer, tf.keras.optimizers.experimental.Optimizer) ), ('OptimizerFactory.build_optimizer returning a non-optimizer object: ' '{}'.format(optimizer)) return optimizer
def build_optimizers_tf_v2(optimizer_config, global_step=None): """Create a TF v2 compatible optimizer based on config. Args: optimizer_config: A Optimizer proto message. global_step: A variable representing the current step. If None, defaults to tf.train.get_or_create_global_step() Returns: An optimizer and a list of variables for summary. Raises: ValueError: when using an unsupported input data type. """ optimizer_type = optimizer_config.WhichOneof('optimizer') optimizer = None summary_vars = [] if optimizer_type == 'rms_prop_optimizer': config = optimizer_config.rms_prop_optimizer learning_rate = _create_learning_rate(config.learning_rate, global_step=global_step) summary_vars.append(learning_rate) optimizer = tf.keras.optimizers.RMSprop( learning_rate, decay=config.decay, momentum=config.momentum_optimizer_value, epsilon=config.epsilon) if optimizer_type == 'momentum_optimizer': config = optimizer_config.momentum_optimizer learning_rate = _create_learning_rate(config.learning_rate, global_step=global_step) summary_vars.append(learning_rate) optimizer = tf.keras.optimizers.SGD( learning_rate, momentum=config.momentum_optimizer_value) if optimizer_type == 'adam_optimizer': config = optimizer_config.adam_optimizer learning_rate = _create_learning_rate(config.learning_rate, global_step=global_step) summary_vars.append(learning_rate) optimizer = tf.keras.optimizers.Adam(learning_rate, epsilon=config.epsilon) if optimizer is None: raise ValueError('Optimizer %s not supported.' % optimizer_type) if optimizer_config.use_moving_average: optimizer = ema_optimizer.ExponentialMovingAverage( optimizer=optimizer, average_decay=optimizer_config.moving_average_decay) return optimizer, summary_vars
def build_optimizer( self, lr: Union[tf.keras.optimizers.schedules.LearningRateSchedule, float], postprocessor: Optional[ Callable[[tf.keras.optimizers.Optimizer], tf.keras.optimizers.Optimizer]] = None): """Build optimizer. Builds optimizer from config. It takes learning rate as input, and builds the optimizer according to the optimizer config. Typically, the learning rate built using self.build_lr() is passed as an argument to this method. Args: lr: A floating point value, or a tf.keras.optimizers.schedules.LearningRateSchedule instance. postprocessor: An optional function for postprocessing the optimizer. It takes an optimizer and returns an optimizer. Returns: tf.keras.optimizers.Optimizer instance. """ optimizer_dict = self._optimizer_config.as_dict() ## Delete clipnorm and clipvalue if None if optimizer_dict['clipnorm'] is None: del optimizer_dict['clipnorm'] if optimizer_dict['clipvalue'] is None: del optimizer_dict['clipvalue'] optimizer_dict['learning_rate'] = lr optimizer = OPTIMIZERS_CLS[self._optimizer_type](**optimizer_dict) if self._use_ema: optimizer = ema_optimizer.ExponentialMovingAverage( optimizer, **self._ema_config.as_dict()) if postprocessor: optimizer = postprocessor(optimizer) assert isinstance(optimizer, tf.keras.optimizers.Optimizer), ( 'OptimizerFactory.build_optimizer returning a non-optimizer object: ' '{}'.format(optimizer)) return optimizer
def add_ema(self, optimizer): """Add EMA to the optimizer independently of the build optimizer method.""" if self._use_ema: optimizer = ema_optimizer.ExponentialMovingAverage( optimizer, **self._ema_config.as_dict()) return optimizer