def _ComputeClassificationLoss(self, predictions, input_batch,
class_weights):
"""Compute classification loss for the given predictions.
Args:
predictions: The output of `ComputePredictions`, contains: logits - [b,
nx, ny, nz, na, 7 + num_classes]. na is the number of anchor
boxes per cell. [..., :7] are (dx, dy, dz, dw, dl, dh, dt).
input_batch: The input batch from which we accesses the groundtruth.
class_weights: Per-class weights to use in loss computation.
Returns:
Classification loss.
"""
p = self.params
predicted_class_logits = py_utils.HasShape(
predictions.classification_logits,
[-1, -1, -1, -1, p.num_anchors, p.num_classes])
bs, nx, ny, nz, na, _ = py_utils.GetShape(predicted_class_logits, 6)
assigned_gt_labels = py_utils.HasShape(input_batch.assigned_gt_labels,
[bs, nx, ny, nz, na])
class_loss = py_utils.SigmoidCrossEntropyFocalLoss(
logits=predicted_class_logits,
labels=tf.one_hot(assigned_gt_labels, p.num_classes),
alpha=p.focal_loss_alpha,
gamma=p.focal_loss_gamma)
class_loss *= class_weights[..., tf.newaxis]
class_loss_sum = tf.reduce_sum(class_loss)
return class_loss_sum
def ComputeLoss(self, theta, predictions, input_batch):
"""Compute loss for the sparse detector model v1.
Args:
theta: A `.NestedMap` object containing variable values of this task.
predictions: A `.NestedMap` object containing residuals and
classification_logits.
input_batch: A `.NestedMap` expected to contain cell_center_xyz,
cell_points_xyz, cell_feature, anchor_bboxes,
anchor_localization_residuals, assigned_gt_labels, and
assigned_cls_mask. See class doc string for details.
Returns:
Two dicts:
- A dict containing str keys and (metric, weight) pairs as values, where
one of the keys is expected to be 'loss'.
- A dict containing arbitrary tensors describing something about each
training example, where the first dimension of each tensor is the batch
index.
"""
p = self.params
batch_size, num_centers = py_utils.GetShape(
input_batch.cell_center_xyz, 2)
# Assert shapes of inputs.
anchor_bboxes = py_utils.HasShape(
input_batch.anchor_bboxes,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 7])
anchor_localization_residuals = py_utils.HasShape(
input_batch.anchor_localization_residuals,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 7])
predicted_residuals = py_utils.HasShape(
predictions.residuals,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 7])
assigned_gt_labels = py_utils.HasShape(
input_batch.assigned_gt_labels,
[batch_size, num_centers, p.num_anchor_bboxes_per_center])
predicted_classification_logits = py_utils.HasShape(
predictions.classification_logits, [
batch_size, num_centers, p.num_anchor_bboxes_per_center,
p.num_classes
])
# assigned_cls_mask is for weighting the classification loss.
# Ignored targets will have their mask = 0; this happens when their IOU is
# not high enough to be a foreground object and not low enough to be
# background.
class_weights = py_utils.HasShape(
input_batch.assigned_cls_mask,
[batch_size, num_centers, p.num_anchor_bboxes_per_center])
class_weights = tf.reshape(
class_weights,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 1])
# Broadcast per class loss weights. For each anchor, there are num_classes
# prediction heads, we weight the outputs of these heads by the per class
# loss weights.
per_class_loss_weight = tf.constant([[[p.per_class_loss_weight]]],
dtype=tf.float32)
per_class_loss_weight = py_utils.HasShape(per_class_loss_weight,
[1, 1, 1, p.num_classes])
class_weights *= per_class_loss_weight
class_weights = py_utils.HasShape(class_weights, [
batch_size, num_centers, p.num_anchor_bboxes_per_center,
p.num_classes
])
# We use assigned_reg_mask for masking the regression loss.
# Only foreground objects will have assigned_reg_mask = 1.
reg_weights = py_utils.HasShape(
input_batch.assigned_reg_mask,
[batch_size, num_centers, p.num_anchor_bboxes_per_center])
reg_weights = tf.reshape(
reg_weights,
[batch_size, num_centers, p.num_anchor_bboxes_per_center, 1])
if p.loss_norm_type == LossNormType.NORM_BY_NUM_POS_PER_CENTER:
# Compute number of positive anchors per example.
foreground_mask = py_utils.HasShape(
input_batch.assigned_reg_mask,
[batch_size, num_centers, p.num_anchor_bboxes_per_center])
# Sum to get the number of foreground anchors for each example.
loss_normalization = tf.reduce_sum(foreground_mask, axis=2)
loss_normalization = tf.maximum(loss_normalization,
tf.ones_like(loss_normalization))
# Reshape for broadcasting.
loss_normalization = tf.reshape(loss_normalization,
[batch_size, num_centers, 1, 1])
# Normalize so that the loss is independent of # centers.
loss_normalization *= num_centers
class_weights /= loss_normalization
reg_weights /= loss_normalization
classification_loss = py_utils.SigmoidCrossEntropyFocalLoss(
logits=predicted_classification_logits,
labels=tf.one_hot(assigned_gt_labels, p.num_classes),
alpha=p.focal_loss_alpha,
gamma=p.focal_loss_gamma)
# Apply mask.
classification_loss *= class_weights
# TODO(jngiam): Consider normalizing by num_foreground_anchors for each
# example instead. This would match the 1/N_positive normalization in
# point pillars.
# Reduce sum over centers, boxes and classes.
classification_loss = tf.reduce_sum(classification_loss,
axis=[1, 2, 3])
# Reduce mean over batch.
classification_loss = tf.reduce_mean(classification_loss)
# Localization regression loss with Huber loss (SmoothL1).
regression_loc_and_dims_loss = self._utils_3d.ScaledHuberLoss(
labels=anchor_localization_residuals[..., :6],
predictions=predicted_residuals[..., :6],
delta=p.huber_loss_delta)
# Rotation loss is computed on a transform on rotation_delta. For a
# direction aware loss, we simply wrap the angles to -pi to pi; for a loss
# that is symmetric to direction (i.e., rotating by pi), we use a sin
# transform.
rotation_delta_transform = tf.sin
if p.direction_aware_rot_loss:
rotation_delta_transform = functools.partial(geometry.WrapAngleRad,
min_val=-np.pi,
max_val=np.pi)
rotation_delta = (predicted_residuals[..., 6:] -
anchor_localization_residuals[..., 6:])
regression_rotation_loss = self._utils_3d.ScaledHuberLoss(
labels=tf.zeros_like(rotation_delta),
predictions=rotation_delta_transform(rotation_delta),
delta=p.huber_loss_delta)
reg_loc_loss = regression_loc_and_dims_loss[..., :3]
reg_dim_loss = regression_loc_and_dims_loss[..., 3:6]
gt_bboxes = self._utils_3d.ResidualsToBBoxes(
anchor_bboxes,
anchor_localization_residuals,
min_angle_rad=-np.pi,
max_angle_rad=np.pi)
predicted_bboxes = self._utils_3d.ResidualsToBBoxes(
anchor_bboxes,
predicted_residuals,
min_angle_rad=-np.pi,
max_angle_rad=np.pi)
# Apply mask to individual losses.
#
# And then reduce sum over centers, boxes, residuals, and batch
# and divide by the batch_size.
regression_rotation_loss *= reg_weights
reg_rot_loss = tf.reduce_sum(regression_rotation_loss) / batch_size
reg_loc_loss *= reg_weights
reg_loc_loss = tf.reduce_sum(reg_loc_loss) / batch_size
reg_dim_loss *= reg_weights
reg_dim_loss = tf.reduce_sum(reg_dim_loss) / batch_size
# Do not create corner loss graph if weight is 0.0
# TODO(bcyang): Remove condition after fixing corner loss NaN issue
if p.corner_loss_weight != 0.0:
reg_corner_loss = self._utils_3d.CornerLoss(
gt_bboxes=gt_bboxes, predicted_bboxes=predicted_bboxes)
reg_corner_loss = tf.expand_dims(reg_corner_loss, axis=-1)
reg_corner_loss *= reg_weights
reg_corner_loss = tf.reduce_sum(reg_corner_loss) / batch_size
else:
reg_corner_loss = 0.0
# Sum components of regression loss.
regression_loss = (p.location_loss_weight * reg_loc_loss +
p.dimension_loss_weight * reg_dim_loss +
p.rotation_loss_weight * reg_rot_loss +
p.corner_loss_weight * reg_corner_loss)
# Compute total loss.
total_loss = (p.loss_weight_localization * regression_loss +
p.loss_weight_classification * classification_loss)
metrics_dict = py_utils.NestedMap({
'loss': (total_loss, batch_size),
'loss/regression': (regression_loss, batch_size),
'loss/regression/loc': (reg_loc_loss, batch_size),
'loss/regression/dim': (reg_dim_loss, batch_size),
'loss/regression/rot': (reg_rot_loss, batch_size),
'loss/regression/corner': (reg_corner_loss, batch_size),
'loss/classification': (classification_loss, batch_size),
})
# Calculate dimension errors
dimension_errors_dict = self._BBoxDimensionErrors(
gt_bboxes, predicted_bboxes, reg_weights)
metrics_dict.update(dimension_errors_dict)
per_example_dict = py_utils.NestedMap({
'residuals': predicted_residuals,
'classification_logits': predicted_classification_logits,
'predicted_bboxes': predicted_bboxes,
'gt_bboxes': gt_bboxes,
'reg_weights': reg_weights,
})
return metrics_dict, per_example_dict
def ComputeLoss(self, theta, predictions, input_batch):
"""Computes loss and other metrics for the given predictions.
Args:
theta: A `.NestedMap` object containing variable values of this task.
predictions: The output of `ComputePredictions`, contains: logits - [b,
nx, ny, nz, na, 7 + num_classes]. na is the number of anchor
boxes per cell. [..., :7] are (dx, dy, dz, dw, dl, dh, dt).
input_batch: The input batch from which we accesses the groundtruth.
Returns:
Two dicts defined as BaseTask.ComputeLoss.
"""
p = self.params
predicted_residuals = py_utils.HasShape(
predictions.residuals, [-1, -1, -1, -1, p.num_anchors, 7])
predicted_class_logits = py_utils.HasShape(
predictions.classification_logits,
[-1, -1, -1, -1, p.num_anchors, p.num_classes])
bs, nx, ny, nz, na, _ = py_utils.GetShape(predicted_class_logits, 6)
# Compute class and regression weights.
class_weights = input_batch.assigned_cls_mask
class_weights = py_utils.HasShape(class_weights, [bs, nx, ny, nz, na])
reg_weights = input_batch.assigned_reg_mask
reg_weights = py_utils.HasShape(reg_weights, [bs, nx, ny, nz, na])
reg_weights = tf.expand_dims(reg_weights, -1)
if p.loss_norm_type == LossNormType.NORM_BY_NUM_POSITIVES:
# Compute number of positive anchors per example.
foreground_mask = py_utils.HasShape(input_batch.assigned_reg_mask,
[bs, nx, ny, nz, na])
# Sum to get the number of foreground anchors for each example.
loss_normalization = tf.reduce_sum(foreground_mask,
axis=[1, 2, 3, 4])
loss_normalization = tf.maximum(loss_normalization,
tf.ones_like(loss_normalization))
# Reshape for broadcasting.
loss_normalization = tf.reshape(loss_normalization,
[bs, 1, 1, 1, 1, 1])
class_weights /= loss_normalization
reg_weights /= loss_normalization
# Classification loss.
assigned_gt_labels = py_utils.HasShape(input_batch.assigned_gt_labels,
[bs, nx, ny, nz, na])
class_loss = py_utils.SigmoidCrossEntropyFocalLoss(
logits=predicted_class_logits,
labels=tf.one_hot(assigned_gt_labels, p.num_classes),
alpha=p.focal_loss_alpha,
gamma=p.focal_loss_gamma)
class_loss *= class_weights[..., tf.newaxis]
class_loss_sum = tf.reduce_sum(class_loss)
# Regression loss.
anchor_localization_residuals = py_utils.HasShape(
input_batch.anchor_localization_residuals, [bs, nx, ny, nz, na, 7])
# Location and dimensions loss.
reg_loc_and_dims_loss = self._utils.ScaledHuberLoss(
predictions=py_utils.HasShape(predicted_residuals[..., :6],
[bs, nx, ny, nz, na, 6]),
labels=anchor_localization_residuals[..., :6],
delta=1 / (3.**2))
# Rotation loss with SmoothL1(sin(delta)).
rot_delta = (predicted_residuals[..., 6:] -
input_batch.anchor_localization_residuals[..., 6:])
if p.use_atan2_heading_loss:
atan2_of_delta = tf.atan2(tf.sin(rot_delta), tf.cos(rot_delta))
reg_rot_loss = self._utils.ScaledHuberLoss(
predictions=atan2_of_delta,
labels=tf.zeros_like(atan2_of_delta),
delta=1 / (3.**2))
else:
# Rotation loss with SmoothL1(sin(delta)).
reg_rot_loss = self._utils.ScaledHuberLoss(
predictions=tf.sin(rot_delta),
labels=tf.zeros_like(rot_delta),
delta=1 / (3.**2))
# Direction loss
if p.direction_classifier_weight > 0.0:
# The target rotations are in the assigned_gt_bbox tensor,
# which already has assigned a gt bounding box to every anchor.
rot_target = input_batch.assigned_gt_bbox[..., 6]
# If rotation is > 0, the class is 1, else it is 0.
rot_dir = tf.cast(rot_target > 0., tf.int32)
# Compute one-hot labels as a target.
rot_dir_onehot = tf.one_hot(rot_dir, 2)
# Manually handle loss reduction.
dir_loss = tf.losses.softmax_cross_entropy(
onehot_labels=rot_dir_onehot,
logits=predictions.predicted_dir,
weights=tf.squeeze(reg_weights, axis=-1),
reduction=tf.losses.Reduction.NONE)
# Reduce across all dimensions (we'll divide by the batch size below).
dir_loss_sum = tf.reduce_sum(dir_loss)
else:
dir_loss_sum = 0.0
# Compute loss contribution from location and dimension separately.
reg_loc_loss = reg_loc_and_dims_loss[..., :3] * reg_weights
reg_loc_loss_sum = tf.reduce_sum(reg_loc_loss)
reg_dim_loss = reg_loc_and_dims_loss[..., 3:6] * reg_weights
reg_dim_loss_sum = tf.reduce_sum(reg_dim_loss)
# Compute rotation loss contribution.
reg_rot_loss *= reg_weights
reg_rot_loss_sum = tf.reduce_sum(reg_rot_loss)
# Num. predictions.
# TODO(zhifengc): Consider other normalization factors. E.g., # of bboxes.
preds = tf.cast(bs, class_loss_sum.dtype)
# Normalize all of the components by batch size.
reg_loc_loss = reg_loc_loss_sum / preds
reg_dim_loss = reg_dim_loss_sum / preds
reg_rot_loss = reg_rot_loss_sum / preds
class_loss = class_loss_sum / preds
dir_loss = dir_loss_sum / preds
# Compute total localization regression loss.
reg_loss = (p.location_loss_weight * reg_loc_loss +
p.dimension_loss_weight * reg_dim_loss +
p.rotation_loss_weight * reg_rot_loss)
# Apply weights to normalized class losses.
loss = (class_loss * p.classification_loss_weight +
reg_loss * p.localization_loss_weight +
dir_loss * p.direction_classifier_weight)
metrics_dict = {
'loss': (loss, preds),
'loss/class': (class_loss, preds),
'loss/reg': (reg_loss, preds),
'loss/reg/rot': (reg_rot_loss, preds),
'loss/reg/loc': (reg_loc_loss, preds),
'loss/reg/dim': (reg_dim_loss, preds),
'loss/dir': (dir_loss, preds),
}
# Calculate dimension errors
min_angle_rad = -np.pi if p.use_atan2_heading_loss else 0
gt_bboxes = self._utils_3d.ResidualsToBBoxes(
input_batch.anchor_bboxes,
anchor_localization_residuals,
min_angle_rad=min_angle_rad,
max_angle_rad=np.pi)
predicted_bboxes = self._utils_3d.ResidualsToBBoxes(
input_batch.anchor_bboxes,
predicted_residuals,
min_angle_rad=min_angle_rad,
max_angle_rad=np.pi)
dimension_errors_dict = self._BBoxDimensionErrors(
gt_bboxes, predicted_bboxes, reg_weights)
metrics_dict.update(dimension_errors_dict)
per_example_dict = {
'residuals': predicted_residuals,
'classification_logits': predicted_class_logits,
}
return metrics_dict, per_example_dict
