def graph_fn(prediction, target):
weights = tf.constant([
[[1.0], [1.0]],
[[1.0], [1.0]],
])
loss = centernet_losses.PenaltyReducedLogisticFocalLoss(alpha=2.0,
beta=0.5)
computed_value = loss(prediction, target, weights=weights)
return computed_value
def build_losses(self, outputs, labels, aux_losses=None):
"""Build losses."""
input_size = self.task_config.model.input_size[0:2]
output_size = outputs['ct_heatmaps'][0].get_shape().as_list()[1:3]
gt_label = tf.map_fn(
# pylint: disable=g-long-lambda
fn=lambda x: target_assigner.assign_centernet_targets(
labels=x,
input_size=input_size,
output_size=output_size,
num_classes=self.task_config.model.num_classes,
max_num_instances=self.task_config.model.max_num_instances,
gaussian_iou=self.task_config.losses.gaussian_iou,
class_offset=self.task_config.losses.class_offset),
elems=labels,
fn_output_signature={
'ct_heatmaps':
tf.TensorSpec(shape=[
output_size[0], output_size[1],
self.task_config.model.num_classes
],
dtype=tf.float32),
'ct_offset':
tf.TensorSpec(
shape=[self.task_config.model.max_num_instances, 2],
dtype=tf.float32),
'size':
tf.TensorSpec(
shape=[self.task_config.model.max_num_instances, 2],
dtype=tf.float32),
'box_mask':
tf.TensorSpec(shape=[self.task_config.model.max_num_instances],
dtype=tf.int32),
'box_indices':
tf.TensorSpec(
shape=[self.task_config.model.max_num_instances, 2],
dtype=tf.int32),
})
losses = {}
# Create loss functions
object_center_loss_fn = centernet_losses.PenaltyReducedLogisticFocalLoss(
)
localization_loss_fn = centernet_losses.L1LocalizationLoss()
# Set up box indices so that they have a batch element as well
box_indices = loss_ops.add_batch_to_indices(gt_label['box_indices'])
box_mask = tf.cast(gt_label['box_mask'], dtype=tf.float32)
num_boxes = tf.cast(loss_ops.get_num_instances_from_weights(
gt_label['box_mask']),
dtype=tf.float32)
# Calculate center heatmap loss
output_unpad_image_shapes = tf.math.ceil(
tf.cast(labels['unpad_image_shapes'], tf.float32) /
self._net_down_scale)
valid_anchor_weights = loss_ops.get_valid_anchor_weights_in_flattened_image(
output_unpad_image_shapes, output_size[0], output_size[1])
valid_anchor_weights = tf.expand_dims(valid_anchor_weights, 2)
pred_ct_heatmap_list = outputs['ct_heatmaps']
true_flattened_ct_heatmap = loss_ops.flatten_spatial_dimensions(
gt_label['ct_heatmaps'])
true_flattened_ct_heatmap = tf.cast(true_flattened_ct_heatmap,
tf.float32)
total_center_loss = 0.0
for ct_heatmap in pred_ct_heatmap_list:
pred_flattened_ct_heatmap = loss_ops.flatten_spatial_dimensions(
ct_heatmap)
pred_flattened_ct_heatmap = tf.cast(pred_flattened_ct_heatmap,
tf.float32)
total_center_loss += object_center_loss_fn(
target_tensor=true_flattened_ct_heatmap,
prediction_tensor=pred_flattened_ct_heatmap,
weights=valid_anchor_weights)
center_loss = tf.reduce_sum(total_center_loss) / float(
len(pred_ct_heatmap_list) * num_boxes)
losses['ct_loss'] = center_loss
# Calculate scale loss
pred_scale_list = outputs['ct_size']
true_scale = tf.cast(gt_label['size'], tf.float32)
total_scale_loss = 0.0
for scale_map in pred_scale_list:
pred_scale = loss_ops.get_batch_predictions_from_indices(
scale_map, box_indices)
pred_scale = tf.cast(pred_scale, tf.float32)
# Only apply loss for boxes that appear in the ground truth
total_scale_loss += tf.reduce_sum(localization_loss_fn(
target_tensor=true_scale, prediction_tensor=pred_scale),
axis=-1) * box_mask
scale_loss = tf.reduce_sum(total_scale_loss) / float(
len(pred_scale_list) * num_boxes)
losses['scale_loss'] = scale_loss
# Calculate offset loss
pred_offset_list = outputs['ct_offset']
true_offset = tf.cast(gt_label['ct_offset'], tf.float32)
total_offset_loss = 0.0
for offset_map in pred_offset_list:
pred_offset = loss_ops.get_batch_predictions_from_indices(
offset_map, box_indices)
pred_offset = tf.cast(pred_offset, tf.float32)
# Only apply loss for boxes that appear in the ground truth
total_offset_loss += tf.reduce_sum(localization_loss_fn(
target_tensor=true_offset, prediction_tensor=pred_offset),
axis=-1) * box_mask
offset_loss = tf.reduce_sum(total_offset_loss) / float(
len(pred_offset_list) * num_boxes)
losses['ct_offset_loss'] = offset_loss
# Aggregate and finalize loss
loss_weights = self.task_config.losses.detection
total_loss = (loss_weights.object_center_weight * center_loss +
loss_weights.scale_weight * scale_loss +
loss_weights.offset_weight * offset_loss)
if aux_losses:
total_loss += tf.add_n(aux_losses)
losses['total_loss'] = total_loss
return losses