This class computes the loss for SparseRCNN. 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)
| 21 | |
| 22 | |
| 23 | class SetCriterion(nn.Module): |
| 24 | """ This class computes the loss for SparseRCNN. |
| 25 | The process happens in two steps: |
| 26 | 1) we compute hungarian assignment between ground truth boxes and the outputs of the model |
| 27 | 2) we supervise each pair of matched ground-truth / prediction (supervise class and box) |
| 28 | """ |
| 29 | def __init__(self, cfg, num_classes, matcher, weight_dict, eos_coef, losses, use_focal): |
| 30 | """ Create the criterion. |
| 31 | Parameters: |
| 32 | num_classes: number of object categories, omitting the special no-object category |
| 33 | matcher: module able to compute a matching between targets and proposals |
| 34 | weight_dict: dict containing as key the names of the losses and as values their relative weight. |
| 35 | eos_coef: relative classification weight applied to the no-object category |
| 36 | losses: list of all the losses to be applied. See get_loss for list of available losses. |
| 37 | """ |
| 38 | super().__init__() |
| 39 | self.cfg = cfg |
| 40 | self.num_classes = num_classes |
| 41 | self.matcher = matcher |
| 42 | self.weight_dict = weight_dict |
| 43 | self.eos_coef = eos_coef |
| 44 | self.losses = losses |
| 45 | self.use_focal = use_focal |
| 46 | if self.use_focal: |
| 47 | self.focal_loss_alpha = cfg.MODEL.SparseRCNN.ALPHA |
| 48 | self.focal_loss_gamma = cfg.MODEL.SparseRCNN.GAMMA |
| 49 | else: |
| 50 | empty_weight = torch.ones(self.num_classes + 1) |
| 51 | empty_weight[-1] = self.eos_coef |
| 52 | self.register_buffer('empty_weight', empty_weight) |
| 53 | |
| 54 | def loss_labels(self, outputs, targets, indices, num_boxes, log=False): |
| 55 | """Classification loss (NLL) |
| 56 | targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes] |
| 57 | """ |
| 58 | assert 'pred_logits' in outputs |
| 59 | src_logits = outputs['pred_logits'] |
| 60 | |
| 61 | idx = self._get_src_permutation_idx(indices) |
| 62 | target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) |
| 63 | target_classes = torch.full(src_logits.shape[:2], self.num_classes, |
| 64 | dtype=torch.int64, device=src_logits.device) |
| 65 | target_classes[idx] = target_classes_o |
| 66 | |
| 67 | if self.use_focal: |
| 68 | src_logits = src_logits.flatten(0, 1) |
| 69 | # prepare one_hot target. |
| 70 | target_classes = target_classes.flatten(0, 1) |
| 71 | pos_inds = torch.nonzero(target_classes != self.num_classes, as_tuple=True)[0] |
| 72 | labels = torch.zeros_like(src_logits) |
| 73 | labels[pos_inds, target_classes[pos_inds]] = 1 |
| 74 | # comp focal loss. |
| 75 | class_loss = sigmoid_focal_loss_jit( |
| 76 | src_logits, |
| 77 | labels, |
| 78 | alpha=self.focal_loss_alpha, |
| 79 | gamma=self.focal_loss_gamma, |
| 80 | reduction="sum", |