Method joint_train_step

official/modeling/multitask/multitask.py:104–149

The joint train step. Args: task_inputs: a dictionary of task names and per-task features. multi_task_model: a MultiTaskBaseModel instance. optimizer: a tf.optimizers.Optimizer. task_metrics: a dictionary of task names and per-task metrics. **kwargs: other argument

(self, task_inputs,
                       multi_task_model: base_model.MultiTaskBaseModel,
                       optimizer: tf_keras.optimizers.Optimizer, task_metrics,
                       **kwargs)

        dp_config=dp_config)

  def joint_train_step(self, task_inputs,
                       multi_task_model: base_model.MultiTaskBaseModel,
                       optimizer: tf_keras.optimizers.Optimizer, task_metrics,
                       **kwargs):
    """The joint train step.

    Args:
      task_inputs: a dictionary of task names and per-task features.
      multi_task_model: a MultiTaskBaseModel instance.
      optimizer: a tf.optimizers.Optimizer.
      task_metrics: a dictionary of task names and per-task metrics.
      **kwargs: other arguments to pass through.

    Returns:
      A dictionary of losses, inculding per-task losses and their weighted sum.
    """
    losses = {}
    with tf.GradientTape() as tape:
      total_loss = 0.0
      for name, model in multi_task_model.sub_tasks.items():
        inputs = task_inputs[name]
        if isinstance(inputs, tuple) and len(inputs) == 2:
          features, labels = inputs
        elif isinstance(inputs, dict):
          features, labels = inputs, inputs
        else:
          raise ValueError("The iterator output is neither a tuple nor a "
                           "dictionary. It is not implemented to support "
                           "such outputs.")
        outputs = model(features, training=True)
        task_loss = self.tasks[name].build_losses(labels, outputs)
        task_weight = self.task_weight(name)
        total_loss += task_weight * task_loss
        losses[name] = task_loss
        self.tasks[name].process_metrics(task_metrics[name], labels, outputs,
                                         **kwargs)

      # Scales loss as the default gradients allreduce performs sum inside
      # the optimizer.
      scaled_loss = total_loss / tf.distribute.get_strategy(
      ).num_replicas_in_sync
    tvars = multi_task_model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    optimizer.apply_gradients(list(zip(grads, tvars)))
    losses["total_loss"] = total_loss
    return losses