This class computes the loss for DETR. The process happens in two steps: 1) we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
| 385 | return "\n".join(lines) |
| 386 | |
| 387 | class DetSetCriterion(nn.Module): |
| 388 | """ This class computes the loss for DETR. |
| 389 | The process happens in two steps: |
| 390 | 1) we compute hungarian assignment between ground truth boxes and the outputs of the model |
| 391 | 2) we supervise each pair of matched ground-truth / prediction (supervise class and box) |
| 392 | """ |
| 393 | def __init__(self, num_classes, matcher, weight_dict, losses, focal_alpha, ign_thr, ginfo): |
| 394 | """ Create the criterion. |
| 395 | Parameters: |
| 396 | num_classes: number of object categories, omitting the special no-object category |
| 397 | matcher: module able to compute a matching between targets and proposals |
| 398 | weight_dict: dict containing as key the names of the losses and as values their relative weight. |
| 399 | losses: list of all the losses to be applied. See get_loss for list of available losses. |
| 400 | focal_alpha: alpha in Focal Loss |
| 401 | """ |
| 402 | super().__init__() |
| 403 | self.num_classes = num_classes |
| 404 | self.matcher = matcher |
| 405 | self.weight_dict = weight_dict |
| 406 | self.losses = losses |
| 407 | self.focal_alpha = focal_alpha |
| 408 | self.ign_thr = ign_thr |
| 409 | self.ginfo = ginfo |
| 410 | |
| 411 | def loss_labels(self, outputs, targets, indices, num_boxes, log=True): |
| 412 | """Classification loss (NLL) |
| 413 | targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes] |
| 414 | """ |
| 415 | assert 'pred_logits' in outputs |
| 416 | src_logits = outputs['pred_logits'] |
| 417 | |
| 418 | idx = self._get_src_permutation_idx(indices) |
| 419 | target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) |
| 420 | target_classes = torch.full(src_logits.shape[:2], self.num_classes, |
| 421 | dtype=torch.int64, device=src_logits.device) |
| 422 | target_classes[idx] = target_classes_o |
| 423 | |
| 424 | with torch.no_grad(): |
| 425 | src_boxes = outputs['pred_boxes'] |
| 426 | valid_mask = torch.stack([ |
| 427 | torch.all(box_ops.box_ioa( |
| 428 | box_ops.box_cxcywh_to_xyxy(boxes), |
| 429 | box_ops.box_cxcywh_to_xyxy(target["ignore"])) < self.ign_thr, 1) |
| 430 | for boxes, target in zip(src_boxes, targets) |
| 431 | ]) | (target_classes != self.num_classes) |
| 432 | if outputs['mask'] is not None: |
| 433 | valid_mask &= outputs['mask']['mask'] |
| 434 | |
| 435 | src_logits = src_logits[valid_mask] |
| 436 | target_classes = target_classes[valid_mask] |
| 437 | |
| 438 | pos_inds = torch.nonzero(target_classes != self.num_classes, as_tuple=True)[0] |
| 439 | labels = torch.zeros_like(src_logits) |
| 440 | labels[pos_inds, target_classes[pos_inds]] = 1 |
| 441 | loss_ce = sigmoid_focal_loss(src_logits, labels, num_boxes, alpha=self.focal_alpha, gamma=2) |
| 442 | losses = {'loss_ce': loss_ce, "valid_ratio": valid_mask.float().mean()} |
| 443 | |
| 444 | if log: |