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)
| 68 | return attention |
| 69 | |
| 70 | class SetCriterion(nn.Module): |
| 71 | """ This class computes the loss for DETR. |
| 72 | The process happens in two steps: |
| 73 | 1) we compute hungarian assignment between ground truth boxes and the outputs of the model |
| 74 | 2) we supervise each pair of matched ground-truth / prediction (supervise class and box) |
| 75 | """ |
| 76 | def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses): |
| 77 | """ Create the criterion. |
| 78 | Parameters: |
| 79 | num_classes: number of object categories, omitting the special no-object category |
| 80 | matcher: module able to compute a matching between targets and proposals |
| 81 | weight_dict: dict containing as key the names of the losses and as values their relative weight. |
| 82 | eos_coef: relative classification weight applied to the no-object category |
| 83 | losses: list of all the losses to be applied. See get_loss for list of available losses. |
| 84 | """ |
| 85 | super().__init__() |
| 86 | self.num_classes = num_classes |
| 87 | self.matcher = matcher |
| 88 | self.weight_dict = weight_dict |
| 89 | self.eos_coef = eos_coef |
| 90 | self.losses = losses |
| 91 | empty_weight = torch.ones(self.num_classes + 1) |
| 92 | empty_weight[-1] = self.eos_coef |
| 93 | self.register_buffer('empty_weight', empty_weight) |
| 94 | |
| 95 | def loss_labels(self, outputs, targets, indices, num_boxes, log=True): |
| 96 | """Classification loss (NLL) |
| 97 | targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes] |
| 98 | """ |
| 99 | assert 'pred_logits' in outputs |
| 100 | src_logits = outputs['pred_logits'] |
| 101 | |
| 102 | idx = self._get_src_permutation_idx(indices) |
| 103 | target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)]) |
| 104 | target_classes = torch.full(src_logits.shape[:2], self.num_classes, |
| 105 | dtype=torch.int64, device=src_logits.device) |
| 106 | target_classes[idx] = target_classes_o |
| 107 | |
| 108 | loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight) |
| 109 | losses = {'loss_ce': loss_ce} |
| 110 | |
| 111 | if log: |
| 112 | # TODO this should probably be a separate loss, not hacked in this one here |
| 113 | losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0] |
| 114 | return losses |
| 115 | |
| 116 | @torch.no_grad() |
| 117 | def loss_cardinality(self, outputs, targets, indices, num_boxes): |
| 118 | """ Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes |
| 119 | This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients |
| 120 | """ |
| 121 | pred_logits = outputs['pred_logits'] |
| 122 | device = pred_logits.device |
| 123 | tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device) |
| 124 | # Count the number of predictions that are NOT "no-object" (which is the last class) |
| 125 | card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1) |
| 126 | card_err = F.l1_loss(card_pred.float(), tgt_lengths.float()) |
| 127 | losses = {'cardinality_error': card_err} |