head.py 文件源码

def _train_op(loss,
              labels,
              train_op_fn,
              centered_bias=None,
              logits_dimension=None,
              loss_fn=None):
  """Returns op for the training step."""
  if centered_bias is not None:
    centered_bias_step = _centered_bias_step(centered_bias, logits_dimension,
                                             labels, loss_fn)
  else:
    centered_bias_step = None
  with ops.name_scope(None, "train_op", (loss, labels)):
    train_op = train_op_fn(loss)
    if centered_bias_step is not None:
      train_op = control_flow_ops.group(train_op, centered_bias_step)
    return train_op