official/vision/losses/maskrcnn_losses.py
# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Losses for maskrcnn model."""
# Import libraries
import tensorflow as tf, tf_keras
class RpnScoreLoss(object):
"""Region Proposal Network score loss function."""
def __init__(self, rpn_batch_size_per_im):
self._rpn_batch_size_per_im = rpn_batch_size_per_im
self._binary_crossentropy = tf_keras.losses.BinaryCrossentropy(
reduction=tf_keras.losses.Reduction.SUM, from_logits=True)
def __call__(self, score_outputs, labels):
"""Computes total RPN detection loss.
Computes total RPN detection loss including box and score from all levels.
Args:
score_outputs: an OrderDict with keys representing levels and values
representing scores in [batch_size, height, width, num_anchors].
labels: the dictionary that returned from dataloader that includes
ground-truth targets.
Returns:
rpn_score_loss: a scalar tensor representing total score loss.
"""
with tf.name_scope('rpn_loss'):
levels = sorted(score_outputs.keys())
score_losses = []
for level in levels:
score_losses.append(
self._rpn_score_loss(
score_outputs[level],
labels[level],
normalizer=tf.cast(
tf.shape(score_outputs[level])[0] *
self._rpn_batch_size_per_im,
dtype=score_outputs[level].dtype)))
# Sums per level losses to total loss.
return tf.math.add_n(score_losses)
def _rpn_score_loss(self, score_outputs, score_targets, normalizer=1.0):
"""Computes score loss."""
# score_targets has three values:
# (1) score_targets[i]=1, the anchor is a positive sample.
# (2) score_targets[i]=0, negative.
# (3) score_targets[i]=-1, the anchor is don't care (ignore).
with tf.name_scope('rpn_score_loss'):
mask = tf.math.logical_or(tf.math.equal(score_targets, 1),
tf.math.equal(score_targets, 0))
score_targets = tf.math.maximum(score_targets,
tf.zeros_like(score_targets))
score_targets = tf.expand_dims(score_targets, axis=-1)
score_outputs = tf.expand_dims(score_outputs, axis=-1)
score_loss = self._binary_crossentropy(
score_targets, score_outputs, sample_weight=mask)
score_loss /= normalizer
return score_loss
class RpnBoxLoss(object):
"""Region Proposal Network box regression loss function."""
def __init__(self, huber_loss_delta: float):
# The delta is typically around the mean value of regression target.
# for instances, the regression targets of 512x512 input with 6 anchors on
# P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].
self._huber_loss = tf_keras.losses.Huber(
delta=huber_loss_delta, reduction=tf_keras.losses.Reduction.SUM)
def __call__(self, box_outputs, labels):
"""Computes total RPN detection loss.
Computes total RPN detection loss including box and score from all levels.
Args:
box_outputs: an OrderDict with keys representing levels and values
representing box regression targets in [batch_size, height, width,
num_anchors * 4].
labels: the dictionary that returned from dataloader that includes
ground-truth targets.
Returns:
rpn_box_loss: a scalar tensor representing total box regression loss.
"""
with tf.name_scope('rpn_loss'):
levels = sorted(box_outputs.keys())
box_losses = []
for level in levels:
box_losses.append(self._rpn_box_loss(box_outputs[level], labels[level]))
# Sum per level losses to total loss.
return tf.add_n(box_losses)
def _rpn_box_loss(self, box_outputs, box_targets, normalizer=1.0):
"""Computes box regression loss."""
with tf.name_scope('rpn_box_loss'):
_, height, width, num_anchors_vertices = box_targets.get_shape().as_list()
# (batch_size, height, width, num_anchors, 4)
reshaped_box_targets = tf.reshape(
box_targets, [-1, height, width, num_anchors_vertices // 4, 4])
# The box is valid if at least one of the ymin, xmin, ymax, ymax is not 0.
# (batch_size, height, width, num_anchors)
valid_mask = tf.reduce_any(
tf.math.abs(reshaped_box_targets) > 1e-6, axis=-1)
# (batch_size, height, width, num_anchors * 4)
valid_mask = tf.cast(
tf.repeat(valid_mask, 4, axis=-1), dtype=box_outputs.dtype)
# (batch_size, height, width, num_anchors * 4, 1)
box_targets = tf.expand_dims(box_targets, axis=-1)
# (batch_size, height, width, num_anchors * 4, 1)
box_outputs = tf.expand_dims(box_outputs, axis=-1)
box_loss = self._huber_loss(
box_targets, box_outputs, sample_weight=valid_mask)
# The loss is normalized by the sum of non-zero weights and additional
# normalizer provided by the function caller. Using + 0.01 here to avoid
# division by zero. For each replica, get the sum of non-zero masks. Then
# get the mean of sums from all replicas. Note there is an extra division
# by `num_replicas` in train_step(). So it is equivalent to normalizing
# the box loss by the global sum of non-zero masks.
replica_context = tf.distribute.get_replica_context()
valid_mask = tf.reduce_sum(valid_mask)
valid_mask_mean = replica_context.all_reduce(
tf.distribute.ReduceOp.MEAN, valid_mask
)
box_loss /= normalizer * (valid_mask_mean + 0.01)
return box_loss
class FastrcnnClassLoss(object):
"""Fast R-CNN classification loss function."""
def __init__(self,
use_binary_cross_entropy: bool = False,
top_k_percent: float = 1.0):
"""Initializes loss computation.
Args:
use_binary_cross_entropy: If true, uses binary cross entropy loss,
otherwise uses categorical cross entropy loss.
top_k_percent: a float, the value lies in [0.0, 1.0]. When its value < 1.,
only aggregate the top k percent of losses. This is useful for hard
example mining.
"""
self._use_binary_cross_entropy = use_binary_cross_entropy
self._top_k_percent = top_k_percent
def __call__(self, class_outputs, class_targets, class_weights=None):
"""Computes the class loss (Fast-RCNN branch) of Mask-RCNN.
This function implements the classification loss of the Fast-RCNN.
The classification loss is categorical (or binary) cross entropy on all
RoIs.
Reference:
https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py
# pylint: disable=line-too-long
Args:
class_outputs: a float tensor representing the class prediction for each
box with a shape of [batch_size, num_boxes, num_classes].
class_targets: a float tensor representing the class label for each box
with a shape of [batch_size, num_boxes].
class_weights: A float list containing the weight of each class.
Returns:
a scalar tensor representing total class loss.
"""
with tf.name_scope('fast_rcnn_loss'):
output_dtype = class_outputs.dtype
num_classes = class_outputs.get_shape().as_list()[-1]
class_weights = (
class_weights if class_weights is not None else [1.0] * num_classes
)
if num_classes != len(class_weights):
raise ValueError(
'Length of class_weights should be {}'.format(num_classes)
)
class_weights = tf.constant(class_weights, dtype=output_dtype)
class_targets_one_hot = tf.one_hot(
tf.cast(class_targets, dtype=tf.int32),
num_classes,
dtype=class_outputs.dtype)
if self._use_binary_cross_entropy:
# (batch_size, num_boxes, num_classes)
cross_entropy_loss = tf.nn.sigmoid_cross_entropy_with_logits(
labels=class_targets_one_hot, logits=class_outputs)
cross_entropy_loss *= class_weights
else:
# (batch_size, num_boxes)
cross_entropy_loss = tf.nn.softmax_cross_entropy_with_logits(
labels=class_targets_one_hot, logits=class_outputs)
class_weight_mask = tf.einsum(
'...y,y->...', class_targets_one_hot, class_weights
)
cross_entropy_loss *= class_weight_mask
if self._top_k_percent < 1.0:
return self.aggregate_loss_top_k(cross_entropy_loss)
else:
return tf.reduce_mean(cross_entropy_loss)
def aggregate_loss_top_k(self, loss, num_valid_values=None):
"""Aggregate the top-k the greatest loss values.
Args:
loss: a float tensor in shape (batch_size, num_boxes) or (batch_size,
num_boxes, num_classes) which stores the loss values.
num_valid_values: the number of loss values which are not ignored. The
default value is None, which means all the loss values are valid.
Returns:
A 0-D float which stores the overall loss of the batch.
"""
loss = tf.reshape(loss, shape=[-1])
top_k_num = tf.cast(
self._top_k_percent * tf.size(loss, out_type=tf.float32), tf.int32)
top_k_losses, _ = tf.math.top_k(loss, k=top_k_num)
normalizer = tf.cast(top_k_num, loss.dtype)
if num_valid_values is not None:
normalizer = tf.minimum(normalizer, tf.cast(num_valid_values, loss.dtype))
return tf.reduce_sum(top_k_losses) / (normalizer + 1e-5)
class FastrcnnBoxLoss(object):
"""Fast R-CNN box regression loss function."""
def __init__(self,
huber_loss_delta: float,
class_agnostic_bbox_pred: bool = False):
"""Initiate Faster RCNN box loss.
Args:
huber_loss_delta: the delta is typically around the mean value of
regression target. For instances, the regression targets of 512x512
input with 6 anchors on P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].
class_agnostic_bbox_pred: if True, class agnostic bounding box prediction
is performed.
"""
self._huber_loss = tf_keras.losses.Huber(
delta=huber_loss_delta, reduction=tf_keras.losses.Reduction.SUM)
self._class_agnostic_bbox_pred = class_agnostic_bbox_pred
def __call__(self, box_outputs, class_targets, box_targets):
"""Computes the box loss (Fast-RCNN branch) of Mask-RCNN.
This function implements the box regression loss of the Fast-RCNN. As the
`box_outputs` produces `num_classes` boxes for each RoI, the reference model
expands `box_targets` to match the shape of `box_outputs` and selects only
the target that the RoI has a maximum overlap. (Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/fast_rcnn.py) # pylint: disable=line-too-long
Instead, this function selects the `box_outputs` by the `class_targets` so
that it doesn't expand `box_targets`.
The box loss is smooth L1-loss on only positive samples of RoIs.
Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py # pylint: disable=line-too-long
Args:
box_outputs: a float tensor representing the box prediction for each box
with a shape of [batch_size, num_boxes, num_classes * 4].
class_targets: a float tensor representing the class label for each box
with a shape of [batch_size, num_boxes].
box_targets: a float tensor representing the box label for each box
with a shape of [batch_size, num_boxes, 4].
Returns:
box_loss: a scalar tensor representing total box regression loss.
"""
with tf.name_scope('fast_rcnn_loss'):
class_targets = tf.cast(class_targets, dtype=tf.int32)
if not self._class_agnostic_bbox_pred:
box_outputs = self._assign_class_targets(box_outputs, class_targets)
return self._fast_rcnn_box_loss(box_outputs, box_targets, class_targets)
def _assign_class_targets(self, box_outputs, class_targets):
"""Selects the box from `box_outputs` based on `class_targets`, with which the box has the maximum overlap."""
_, num_rois, num_class_specific_boxes = box_outputs.get_shape().as_list()
num_classes = num_class_specific_boxes // 4
box_outputs = tf.reshape(box_outputs, [-1, num_rois, num_classes, 4])
class_targets_ont_hot = tf.one_hot(
class_targets, num_classes, dtype=box_outputs.dtype
)
return tf.einsum('bnij,bni->bnj', box_outputs, class_targets_ont_hot)
def _fast_rcnn_box_loss(self, box_outputs, box_targets, class_targets,
normalizer=1.0):
"""Computes box regression loss."""
with tf.name_scope('fast_rcnn_box_loss'):
mask = tf.tile(
tf.expand_dims(tf.greater(class_targets, 0), axis=2), [1, 1, 4])
mask = tf.cast(mask, dtype=box_outputs.dtype)
box_targets = tf.expand_dims(box_targets, axis=-1)
box_outputs = tf.expand_dims(box_outputs, axis=-1)
box_loss = self._huber_loss(box_targets, box_outputs, sample_weight=mask)
# The loss is normalized by the number of ones in mask,
# additional normalizer provided by the user and using 0.01 here to avoid
# division by 0. For each replica, get the sum of non-zero masks. Then
# get the mean of sums from all replicas. Note there is an extra division
# by `num_replicas` in train_step(). So it is equivalent to normalizing
# the box loss by the global sum of non-zero masks.
replica_context = tf.distribute.get_replica_context()
mask = tf.reduce_sum(mask)
mask_mean = replica_context.all_reduce(
tf.distribute.ReduceOp.MEAN, mask
)
box_loss /= normalizer * (mask_mean + 0.01)
return box_loss
class MaskrcnnLoss(object):
"""Mask R-CNN instance segmentation mask loss function."""
def __init__(self):
self._binary_crossentropy = tf_keras.losses.BinaryCrossentropy(
reduction=tf_keras.losses.Reduction.SUM, from_logits=True)
def __call__(self, mask_outputs, mask_targets, select_class_targets):
"""Computes the mask loss of Mask-RCNN.
This function implements the mask loss of Mask-RCNN. As the `mask_outputs`
produces `num_classes` masks for each RoI, the reference model expands
`mask_targets` to match the shape of `mask_outputs` and selects only the
target that the RoI has a maximum overlap. (Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py) # pylint: disable=line-too-long
Instead, this implementation selects the `mask_outputs` by the
`class_targets` so that it doesn't expand `mask_targets`. Note that the
selection logic is done in the post-processing of mask_rcnn_fn in
mask_rcnn_architecture.py.
Args:
mask_outputs: a float tensor representing the prediction for each mask,
with a shape of
[batch_size, num_masks, mask_height, mask_width].
mask_targets: a float tensor representing the binary mask of ground truth
labels for each mask with a shape of
[batch_size, num_masks, mask_height, mask_width].
select_class_targets: a tensor with a shape of [batch_size, num_masks],
representing the foreground mask targets.
Returns:
mask_loss: a float tensor representing total mask loss.
"""
with tf.name_scope('mask_rcnn_loss'):
_, _, mask_height, mask_width = mask_outputs.get_shape().as_list()
weights = tf.tile(
tf.greater(select_class_targets, 0)[:, :, tf.newaxis, tf.newaxis],
[1, 1, mask_height, mask_width],
)
weights = tf.cast(weights, dtype=mask_outputs.dtype)
mask_targets = tf.expand_dims(mask_targets, axis=-1)
mask_outputs = tf.expand_dims(mask_outputs, axis=-1)
mask_loss = self._binary_crossentropy(mask_targets, mask_outputs,
sample_weight=weights)
# For each replica, get the sum of non-zero weights. Then get the mean of
# sums from all replicas. Note there is an extra division by
# `num_replicas` in train_step(). So it is equivalent to normalizing the
# mask loss by the global sum of non-zero weights.
replica_context = tf.distribute.get_replica_context()
weights = tf.reduce_sum(weights)
weights_mean = replica_context.all_reduce(
tf.distribute.ReduceOp.MEAN, weights
)
# The loss is normalized by the number of 1's in weights and
# + 0.01 is used to avoid division by zero.
return mask_loss / (weights_mean + 0.01)