def optimize_fn(model,
                optimizer=None,
                learning_rate_schedule=None,
                loss=None,
                metrics=None):
  """Compile the model in Keras."""
  learning_rate_schedule = learning_rate_schedule or T2TLearningRateSchedule()
  if optimizer:
    optimizer = optimizer(learning_rate=learning_rate_schedule)
  else:  # We use Adam by default with adjusted parameters.
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=learning_rate_schedule,
        beta_1=0.9, beta_2=0.997, epsilon=1e-9)
  metrics = metrics or [tf.keras.metrics.sparse_categorical_accuracy]
  def xent_loss(y, x):
    return tf.keras.backend.sparse_categorical_crossentropy(
        y, x, from_logits=True)
  loss = loss or xent_loss
  return model.compile(optimizer=optimizer,
                       loss=loss,
                       metrics=metrics)
def optimize_fn ( model , optimizer = None , learning_rate_schedule = None , loss = None , metrics = None ) : learning_rate_schedule = learning_rate_schedule or T2TLearningRateSchedule ( ) if optimizer : optimizer = optimizer ( learning_rate = learning_rate_schedule ) else : optimizer = tf . keras . optimizers . Adam ( learning_rate = learning_rate_schedule , beta_1 = 0.9 , beta_2 = 0.997 , epsilon = 1e-9 ) metrics = metrics or [ tf . keras . metrics . sparse_categorical_accuracy ] def xent_loss ( y , x ) : return tf . keras . backend . sparse_categorical_crossentropy ( y , x , from_logits = True ) loss = loss or xent_loss return model . compile ( optimizer = optimizer , loss = loss , metrics = metrics )
Compile the model in Keras.