head.py 文件源码

def _mean_squared_loss(logits, labels):
  # To prevent broadcasting inside "-".
  if len(labels.get_shape()) == 1:
    labels = array_ops.expand_dims(labels, dim=[1])
  # TODO(zakaria): make sure it does not recreate the broadcast bug.
  if len(logits.get_shape()) == 1:
    logits = array_ops.expand_dims(logits, dim=[1])
  logits.get_shape().assert_is_compatible_with(labels.get_shape())
  return math_ops.square(logits - math_ops.to_float(labels))