Function get_trainer

official/modeling/multitask/train_lib.py:153–182

Creates a multi-task trainer for the given task. Args: distribution_strategy: A distribution strategy. params: ExperimentConfig instance. task: A MultiTaskTask instance. model: A MultiTaskBaseModel instance. Returns: An Orbit trainer instance.

(
    distribution_strategy: tf.distribute.Strategy,
    params: configs.MultiEvalExperimentConfig,
    task: multitask.MultiTask,
    model: base_model.MultiTaskBaseModel | tf_keras.Model,
)

def get_trainer(
    distribution_strategy: tf.distribute.Strategy,
    params: configs.MultiEvalExperimentConfig,
    task: multitask.MultiTask,
    model: base_model.MultiTaskBaseModel | tf_keras.Model,
) -> orbit.StandardTrainer:
  """Creates a multi-task trainer for the given task.

  Args:
    distribution_strategy: A distribution strategy.
    params: ExperimentConfig instance.
    task: A MultiTaskTask instance.
    model: A MultiTaskBaseModel instance.

  Returns:
      An Orbit trainer instance.
  """
  with distribution_strategy.scope():
    kwargs = dict(
        multi_task=task,
        multi_task_model=model,
        optimizer=train_utils.create_optimizer(task, params),
    )
    if params.trainer.trainer_type == 'interleaving':
      kwargs.update(
          task_sampler=task_sampler.get_task_sampler(
              params.trainer.task_sampler, task.task_weights
          )
      )
    return TRAINERS[params.trainer.trainer_type](**kwargs)


TrainActionsFactoryType = Callable[