Train the given model on the given dataset.

  Args:
    data_dir: Directory where the data is located.
    output_dir: Directory where to put the logs and checkpoints.
    model_class: The model class to train.
    dataset: The name of the dataset to train on.
    input_names: List of strings with the names of the features on input.
    target_names: List of strings with the names of the target features.
    train_steps: for how many steps to train.
    eval_steps: for how many steps to do evaluation.
    eval_frequency: how often (every this many steps) to run evaluation.

def train_fn(data_dir=None, output_dir=None,
             model_class=gin.REQUIRED, dataset=gin.REQUIRED,
             input_names=None, target_names=None,
             train_steps=1000, eval_steps=1, eval_frequency=100):
  """Train the given model on the given dataset.

  Args:
    data_dir: Directory where the data is located.
    output_dir: Directory where to put the logs and checkpoints.
    model_class: The model class to train.
    dataset: The name of the dataset to train on.
    input_names: List of strings with the names of the features on input.
    target_names: List of strings with the names of the target features.
    train_steps: for how many steps to train.
    eval_steps: for how many steps to do evaluation.
    eval_frequency: how often (every this many steps) to run evaluation.
  """
  train_data, eval_data, features_info, keys = train_and_eval_dataset(
      dataset, data_dir)
  if input_names is None:
    input_names = keys[0]
  if target_names is None:
    target_names = keys[1]
  # TODO(lukaszkaiser): The use of distribution strategy below fails like this:
  #   .../keras/models.py", line 93, in _clone_functional_model
  #      for layer in model._input_layers:
  #   AttributeError: 'BasicFcRelu' object has no attribute '_input_layers'
  # strategy = tf.distribute.MirroredStrategy()
  # with strategy.scope():
  model = model_class(features_info=features_info,
                      input_names=input_names, target_names=target_names)
  optimize_fn(model)
  train_batches = shuffle_and_batch_data(
      train_data, target_names, features_info, training=True)
  eval_batches = shuffle_and_batch_data(
      eval_data, target_names, features_info, training=False)
  # Need to run one training step just to get optimizer variables to load.
  model.fit(train_batches, epochs=1, steps_per_epoch=1)

  # Training loop.
  callbacks = []
  callbacks.append(tf.keras.callbacks.History())
  callbacks.append(tf.keras.callbacks.BaseLogger())
  last_epoch = 0
  if output_dir is not None:
    callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=output_dir))
    output_format = os.path.join(output_dir, "model-{epoch:05d}")
    callbacks.append(tf.keras.callbacks.ModelCheckpoint(
        filepath=output_format, save_weights_only=True))
    checkpoints = tf.gfile.Glob(os.path.join(output_dir, "model-*"))
    # Take basenames and strip the "model-" prefix.
    checkpoints = [os.path.basename(ckpt)[6:] for ckpt in checkpoints]
    # Get epoch numbers from the filenames and sort to obtain last epoch.
    epoch_numbers = [int(ckpt[:5]) for ckpt in checkpoints if len(ckpt) > 4]
    epoch_numbers.sort()
    if epoch_numbers:
      last_epoch = epoch_numbers[-1]
      saved_path = os.path.join(output_dir, "model-%05d" % last_epoch)
      model.load_weights(saved_path)
  model.fit(train_batches,
            epochs=train_steps // eval_frequency,
            steps_per_epoch=eval_frequency,
            validation_data=eval_batches,
            validation_steps=eval_steps,
            initial_epoch=last_epoch,
            callbacks=callbacks)