File size: 12,378 Bytes
f53b39e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from collections import defaultdict
from typing import Dict, List
import torch
import torch.distributed
import torch.nn as nn
import torch.nn.functional as F
from training.trainer import CORE_LOSS_KEY
from training.utils.distributed import get_world_size, is_dist_avail_and_initialized
def dice_loss(inputs, targets, num_objects, loss_on_multimask=False):
"""
Compute the DICE loss, similar to generalized IOU for masks
Args:
inputs: A float tensor of arbitrary shape.
The predictions for each example.
targets: A float tensor with the same shape as inputs. Stores the binary
classification label for each element in inputs
(0 for the negative class and 1 for the positive class).
num_objects: Number of objects in the batch
loss_on_multimask: True if multimask prediction is enabled
Returns:
Dice loss tensor
"""
inputs = inputs.sigmoid()
if loss_on_multimask:
# inputs and targets are [N, M, H, W] where M corresponds to multiple predicted masks
assert inputs.dim() == 4 and targets.dim() == 4
# flatten spatial dimension while keeping multimask channel dimension
inputs = inputs.flatten(2)
targets = targets.flatten(2)
numerator = 2 * (inputs * targets).sum(-1)
else:
inputs = inputs.flatten(1)
numerator = 2 * (inputs * targets).sum(1)
denominator = inputs.sum(-1) + targets.sum(-1)
loss = 1 - (numerator + 1) / (denominator + 1)
if loss_on_multimask:
return loss / num_objects
return loss.sum() / num_objects
def sigmoid_focal_loss(
inputs,
targets,
num_objects,
alpha: float = 0.25,
gamma: float = 2,
loss_on_multimask=False,
):
"""
Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
Args:
inputs: A float tensor of arbitrary shape.
The predictions for each example.
targets: A float tensor with the same shape as inputs. Stores the binary
classification label for each element in inputs
(0 for the negative class and 1 for the positive class).
num_objects: Number of objects in the batch
alpha: (optional) Weighting factor in range (0,1) to balance
positive vs negative examples. Default = -1 (no weighting).
gamma: Exponent of the modulating factor (1 - p_t) to
balance easy vs hard examples.
loss_on_multimask: True if multimask prediction is enabled
Returns:
focal loss tensor
"""
prob = inputs.sigmoid()
ce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
p_t = prob * targets + (1 - prob) * (1 - targets)
loss = ce_loss * ((1 - p_t) ** gamma)
if alpha >= 0:
alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
loss = alpha_t * loss
if loss_on_multimask:
# loss is [N, M, H, W] where M corresponds to multiple predicted masks
assert loss.dim() == 4
return loss.flatten(2).mean(-1) / num_objects # average over spatial dims
return loss.mean(1).sum() / num_objects
def iou_loss(
inputs, targets, pred_ious, num_objects, loss_on_multimask=False, use_l1_loss=False
):
"""
Args:
inputs: A float tensor of arbitrary shape.
The predictions for each example.
targets: A float tensor with the same shape as inputs. Stores the binary
classification label for each element in inputs
(0 for the negative class and 1 for the positive class).
pred_ious: A float tensor containing the predicted IoUs scores per mask
num_objects: Number of objects in the batch
loss_on_multimask: True if multimask prediction is enabled
use_l1_loss: Whether to use L1 loss is used instead of MSE loss
Returns:
IoU loss tensor
"""
assert inputs.dim() == 4 and targets.dim() == 4
pred_mask = inputs.flatten(2) > 0
gt_mask = targets.flatten(2) > 0
area_i = torch.sum(pred_mask & gt_mask, dim=-1).float()
area_u = torch.sum(pred_mask | gt_mask, dim=-1).float()
actual_ious = area_i / torch.clamp(area_u, min=1.0)
if use_l1_loss:
loss = F.l1_loss(pred_ious, actual_ious, reduction="none")
else:
loss = F.mse_loss(pred_ious, actual_ious, reduction="none")
if loss_on_multimask:
return loss / num_objects
return loss.sum() / num_objects
class MultiStepMultiMasksAndIous(nn.Module):
def __init__(
self,
weight_dict,
focal_alpha=0.25,
focal_gamma=2,
supervise_all_iou=False,
iou_use_l1_loss=False,
pred_obj_scores=False,
focal_gamma_obj_score=0.0,
focal_alpha_obj_score=-1,
):
"""
This class computes the multi-step multi-mask and IoU losses.
Args:
weight_dict: dict containing weights for focal, dice, iou losses
focal_alpha: alpha for sigmoid focal loss
focal_gamma: gamma for sigmoid focal loss
supervise_all_iou: if True, back-prop iou losses for all predicted masks
iou_use_l1_loss: use L1 loss instead of MSE loss for iou
pred_obj_scores: if True, compute loss for object scores
focal_gamma_obj_score: gamma for sigmoid focal loss on object scores
focal_alpha_obj_score: alpha for sigmoid focal loss on object scores
"""
super().__init__()
self.weight_dict = weight_dict
self.focal_alpha = focal_alpha
self.focal_gamma = focal_gamma
assert "loss_mask" in self.weight_dict
assert "loss_dice" in self.weight_dict
assert "loss_iou" in self.weight_dict
if "loss_class" not in self.weight_dict:
self.weight_dict["loss_class"] = 0.0
self.focal_alpha_obj_score = focal_alpha_obj_score
self.focal_gamma_obj_score = focal_gamma_obj_score
self.supervise_all_iou = supervise_all_iou
self.iou_use_l1_loss = iou_use_l1_loss
self.pred_obj_scores = pred_obj_scores
def forward(self, outs_batch: List[Dict], targets_batch: torch.Tensor):
assert len(outs_batch) == len(targets_batch)
num_objects = torch.tensor(
(targets_batch.shape[1]), device=targets_batch.device, dtype=torch.float
) # Number of objects is fixed within a batch
if is_dist_avail_and_initialized():
torch.distributed.all_reduce(num_objects)
num_objects = torch.clamp(num_objects / get_world_size(), min=1).item()
losses = defaultdict(int)
for outs, targets in zip(outs_batch, targets_batch):
cur_losses = self._forward(outs, targets, num_objects)
for k, v in cur_losses.items():
losses[k] += v
return losses
def _forward(self, outputs: Dict, targets: torch.Tensor, num_objects):
"""
Compute the losses related to the masks: the focal loss and the dice loss.
and also the MAE or MSE loss between predicted IoUs and actual IoUs.
Here "multistep_pred_multimasks_high_res" is a list of multimasks (tensors
of shape [N, M, H, W], where M could be 1 or larger, corresponding to
one or multiple predicted masks from a click.
We back-propagate focal, dice losses only on the prediction channel
with the lowest focal+dice loss between predicted mask and ground-truth.
If `supervise_all_iou` is True, we backpropagate ious losses for all predicted masks.
"""
target_masks = targets.unsqueeze(1).float()
assert target_masks.dim() == 4 # [N, 1, H, W]
src_masks_list = outputs["multistep_pred_multimasks_high_res"]
ious_list = outputs["multistep_pred_ious"]
object_score_logits_list = outputs["multistep_object_score_logits"]
assert len(src_masks_list) == len(ious_list)
assert len(object_score_logits_list) == len(ious_list)
# accumulate the loss over prediction steps
losses = {"loss_mask": 0, "loss_dice": 0, "loss_iou": 0, "loss_class": 0}
for src_masks, ious, object_score_logits in zip(
src_masks_list, ious_list, object_score_logits_list
):
self._update_losses(
losses, src_masks, target_masks, ious, num_objects, object_score_logits
)
losses[CORE_LOSS_KEY] = self.reduce_loss(losses)
return losses
def _update_losses(
self, losses, src_masks, target_masks, ious, num_objects, object_score_logits
):
target_masks = target_masks.expand_as(src_masks)
# get focal, dice and iou loss on all output masks in a prediction step
loss_multimask = sigmoid_focal_loss(
src_masks,
target_masks,
num_objects,
alpha=self.focal_alpha,
gamma=self.focal_gamma,
loss_on_multimask=True,
)
loss_multidice = dice_loss(
src_masks, target_masks, num_objects, loss_on_multimask=True
)
if not self.pred_obj_scores:
loss_class = torch.tensor(
0.0, dtype=loss_multimask.dtype, device=loss_multimask.device
)
target_obj = torch.ones(
loss_multimask.shape[0],
1,
dtype=loss_multimask.dtype,
device=loss_multimask.device,
)
else:
target_obj = torch.any((target_masks[:, 0] > 0).flatten(1), dim=-1)[
..., None
].float()
loss_class = sigmoid_focal_loss(
object_score_logits,
target_obj,
num_objects,
alpha=self.focal_alpha_obj_score,
gamma=self.focal_gamma_obj_score,
)
loss_multiiou = iou_loss(
src_masks,
target_masks,
ious,
num_objects,
loss_on_multimask=True,
use_l1_loss=self.iou_use_l1_loss,
)
assert loss_multimask.dim() == 2
assert loss_multidice.dim() == 2
assert loss_multiiou.dim() == 2
if loss_multimask.size(1) > 1:
# take the mask indices with the smallest focal + dice loss for back propagation
loss_combo = (
loss_multimask * self.weight_dict["loss_mask"]
+ loss_multidice * self.weight_dict["loss_dice"]
)
best_loss_inds = torch.argmin(loss_combo, dim=-1)
batch_inds = torch.arange(loss_combo.size(0), device=loss_combo.device)
loss_mask = loss_multimask[batch_inds, best_loss_inds].unsqueeze(1)
loss_dice = loss_multidice[batch_inds, best_loss_inds].unsqueeze(1)
# calculate the iou prediction and slot losses only in the index
# with the minimum loss for each mask (to be consistent w/ SAM)
if self.supervise_all_iou:
loss_iou = loss_multiiou.mean(dim=-1).unsqueeze(1)
else:
loss_iou = loss_multiiou[batch_inds, best_loss_inds].unsqueeze(1)
else:
loss_mask = loss_multimask
loss_dice = loss_multidice
loss_iou = loss_multiiou
# backprop focal, dice and iou loss only if obj present
loss_mask = loss_mask * target_obj
loss_dice = loss_dice * target_obj
loss_iou = loss_iou * target_obj
# sum over batch dimension (note that the losses are already divided by num_objects)
losses["loss_mask"] += loss_mask.sum()
losses["loss_dice"] += loss_dice.sum()
losses["loss_iou"] += loss_iou.sum()
losses["loss_class"] += loss_class
def reduce_loss(self, losses):
reduced_loss = 0.0
for loss_key, weight in self.weight_dict.items():
if loss_key not in losses:
raise ValueError(f"{type(self)} doesn't compute {loss_key}")
if weight != 0:
reduced_loss += losses[loss_key] * weight
return reduced_loss
|