def _compute_kp_offset_loss(self, input_height, input_width, task_name, offset_predictions, localization_loss_fn):
    'Computes the offset loss of the keypoint estimation task.\n\n    Args:\n      input_height: An integer scalar tensor representing input image height.\n      input_width: An integer scalar tensor representing input image width.\n      task_name: A string representing the name of the keypoint task.\n      offset_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, 2] representing the prediction heads of the model\n        for keypoint offset.\n      localization_loss_fn: An object_detection.core.losses.Loss object to\n        compute the loss for the keypoint offset predictions in CenterNet.\n\n    Returns:\n      loss: A float scalar tensor representing the keypoint offset loss\n        normalized by number of total keypoints.\n    '
    gt_keypoints_list = self.groundtruth_lists(fields.BoxListFields.keypoints)
    gt_classes_list = self.groundtruth_lists(fields.BoxListFields.classes)
    gt_weights_list = self.groundtruth_lists(fields.BoxListFields.weights)
    assigner = self._target_assigner_dict[task_name]
    (batch_indices, batch_offsets, batch_weights) = assigner.assign_keypoints_offset_targets(height=input_height, width=input_width, gt_keypoints_list=gt_keypoints_list, gt_weights_list=gt_weights_list, gt_classes_list=gt_classes_list)
    loss = 0.0
    for prediction in offset_predictions:
        (batch_size, out_height, out_width, channels) = _get_shape(prediction, 4)
        if (channels > 2):
            prediction = tf.reshape(prediction, shape=[batch_size, out_height, out_width, (- 1), 2])
        prediction = cn_assigner.get_batch_predictions_from_indices(prediction, batch_indices)
        unweighted_loss = localization_loss_fn(prediction, batch_offsets, weights=tf.expand_dims(tf.ones_like(batch_weights), (- 1)))
        loss += (batch_weights * tf.reduce_sum(unweighted_loss, axis=1))
    loss = (tf.reduce_sum(loss) / (float(len(offset_predictions)) * tf.maximum(tf.reduce_sum(batch_weights), 1.0)))
    return loss
Computes the offset loss of the keypoint estimation task.

Args:
  input_height: An integer scalar tensor representing input image height.
  input_width: An integer scalar tensor representing input image width.
  task_name: A string representing the name of the keypoint task.
  offset_predictions: A list of float tensors of shape [batch_size,
    out_height, out_width, 2] representing the prediction heads of the model
    for keypoint offset.
  localization_loss_fn: An object_detection.core.losses.Loss object to
    compute the loss for the keypoint offset predictions in CenterNet.

Returns:
  loss: A float scalar tensor representing the keypoint offset loss
    normalized by number of total keypoints.
def _compute_kp_offset_loss(self, input_height, input_width, task_name, offset_predictions, localization_loss_fn):
    'Computes the offset loss of the keypoint estimation task.\n\n    Args:\n      input_height: An integer scalar tensor representing input image height.\n      input_width: An integer scalar tensor representing input image width.\n      task_name: A string representing the name of the keypoint task.\n      offset_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, 2] representing the prediction heads of the model\n        for keypoint offset.\n      localization_loss_fn: An object_detection.core.losses.Loss object to\n        compute the loss for the keypoint offset predictions in CenterNet.\n\n    Returns:\n      loss: A float scalar tensor representing the keypoint offset loss\n        normalized by number of total keypoints.\n    '
    gt_keypoints_list = self.groundtruth_lists(fields.BoxListFields.keypoints)
    gt_classes_list = self.groundtruth_lists(fields.BoxListFields.classes)
    gt_weights_list = self.groundtruth_lists(fields.BoxListFields.weights)
    assigner = self._target_assigner_dict[task_name]
    (batch_indices, batch_offsets, batch_weights) = assigner.assign_keypoints_offset_targets(height=input_height, width=input_width, gt_keypoints_list=gt_keypoints_list, gt_weights_list=gt_weights_list, gt_classes_list=gt_classes_list)
    loss = 0.0
    for prediction in offset_predictions:
        (batch_size, out_height, out_width, channels) = _get_shape(prediction, 4)
        if (channels > 2):
            prediction = tf.reshape(prediction, shape=[batch_size, out_height, out_width, (- 1), 2])
        prediction = cn_assigner.get_batch_predictions_from_indices(prediction, batch_indices)
        unweighted_loss = localization_loss_fn(prediction, batch_offsets, weights=tf.expand_dims(tf.ones_like(batch_weights), (- 1)))
        loss += (batch_weights * tf.reduce_sum(unweighted_loss, axis=1))
    loss = (tf.reduce_sum(loss) / (float(len(offset_predictions)) * tf.maximum(tf.reduce_sum(batch_weights), 1.0)))
    return loss