Add focal gamma param

cfgs/train.yaml

@@ -81,6 +81,7 @@ giou_loss_coef: 2
 eos_coef: 0.1
 focal_loss: false
 focal_alpha: 0.25
+focal_gamma: 2
 # Dataset
 dataset: coco
 train_split: train

cfgs/train_focal_loss.yaml

@@ -1,4 +1,5 @@
 focal_loss: true
 focal_alpha: 0.25
+focal_gamma: 2
 cls_loss_coef: 2.0
 set_cost_class: 2.0

src/trackformer/engine.py

@@ -226,6 +226,7 @@ def evaluate(model, criterion, postprocessors, data_loader, device,
             results_orig = {
                 target['image_id'].item(): output
                 for target, output in zip(targets, results_orig)}
+
             coco_evaluator.update(results_orig)
 
         if panoptic_evaluator is not None:
@@ -333,6 +334,7 @@ def evaluate(model, criterion, postprocessors, data_loader, device,
     if visualizers:
         vis_epoch = visualizers['epoch_metrics']
         y_data = [stats[legend_name] for legend_name in vis_epoch.viz_opts['legend']]
+
         vis_epoch.plot(y_data, epoch)
 
         visualizers['epoch_eval'].plot(eval_stats, epoch)

src/trackformer/models/__init__.py

@@ -108,6 +108,7 @@ def build_model(args):
         losses=losses,
         focal_loss=args.focal_loss,
         focal_alpha=args.focal_alpha,
+        focal_gamma=args.focal_gamma,
         tracking=args.tracking,
         track_query_false_positive_eos_weight=args.track_query_false_positive_eos_weight,)
     criterion.to(device)

src/trackformer/models/deformable_detr.py

@@ -18,9 +18,7 @@
 from torch import nn
 
 from ..util import box_ops
-from ..util.misc import (NestedTensor, accuracy, get_world_size,
-                         inverse_sigmoid, is_dist_avail_and_initialized,
-                         nested_tensor_from_tensor_list, sigmoid_focal_loss)
+from ..util.misc import NestedTensor, inverse_sigmoid, nested_tensor_from_tensor_list
 from .detr import DETR, PostProcess, SetCriterion
 
 

src/trackformer/models/detr.py

@@ -143,7 +143,7 @@ class SetCriterion(nn.Module):
         2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
     """
     def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses,
-                 focal_loss, focal_alpha, tracking, track_query_false_positive_eos_weight):
+                 focal_loss, focal_alpha, focal_gamma, tracking, track_query_false_positive_eos_weight):
         """ Create the criterion.
         Parameters:
             num_classes: number of object categories, omitting the special no-object category
@@ -165,6 +165,7 @@ def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses,
         self.register_buffer('empty_weight', empty_weight)
         self.focal_loss = focal_loss
         self.focal_alpha = focal_alpha
+        self.focal_gamma = focal_gamma
         self.tracking = tracking
         self.track_query_false_positive_eos_weight = track_query_false_positive_eos_weight
 
@@ -233,7 +234,9 @@ def loss_labels_focal(self, outputs, targets, indices, num_boxes, log=True):
             query_mask = torch.stack([~t['track_queries_placeholder_mask'] for t in targets])[..., None]
             query_mask = query_mask.repeat(1, 1, self.num_classes)
 
-        loss_ce = sigmoid_focal_loss(src_logits, target_classes_onehot, num_boxes, alpha=self.focal_alpha, gamma=2, query_mask=query_mask)
+        loss_ce = sigmoid_focal_loss(
+            src_logits, target_classes_onehot, num_boxes,
+            alpha=self.focal_alpha, gamma=self.focal_gamma, query_mask=query_mask)
 
         if self.tracking:
             mean_num_queries = torch.tensor([len(m.nonzero()) for m in query_mask]).float().mean()
@@ -245,6 +248,26 @@ def loss_labels_focal(self, outputs, targets, indices, num_boxes, log=True):
         if log:
             # TODO this should probably be a separate loss, not hacked in this one here
             losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]
+
+        # compute seperate track and object query losses
+        # loss_ce = sigmoid_focal_loss(
+        #     src_logits, target_classes_onehot, num_boxes,
+        #     alpha=self.focal_alpha, gamma=self.focal_gamma, query_mask=query_mask, reduction=False)
+        # loss_ce *= src_logits.shape[1]
+
+        # track_query_target_masks = []
+        # for t, ind in zip(targets, indices):
+        #     track_query_target_mask = torch.zeros_like(ind[1]).bool()
+
+        #     for i in t['track_query_match_ids']:
+        #         track_query_target_mask[ind[1].eq(i).nonzero()[0]] = True
+
+        #     track_query_target_masks.append(track_query_target_mask)
+        # track_query_target_masks = torch.cat(track_query_target_masks)
+
+        # losses['loss_ce_track_queries'] = loss_ce[idx][track_query_target_masks].mean(1).sum() / num_boxes
+        # losses['loss_ce_object_queries'] = loss_ce[idx][~track_query_target_masks].mean(1).sum() / num_boxes
+
         return losses
 
     @torch.no_grad()
@@ -282,6 +305,24 @@ def loss_boxes(self, outputs, targets, indices, num_boxes):
             box_ops.box_cxcywh_to_xyxy(src_boxes),
             box_ops.box_cxcywh_to_xyxy(target_boxes)))
         losses['loss_giou'] = loss_giou.sum() / num_boxes
+
+        # compute seperate track and object query losses
+        # track_query_target_masks = []
+        # for t, ind in zip(targets, indices):
+        #     track_query_target_mask = torch.zeros_like(ind[1]).bool()
+
+        #     for i in t['track_query_match_ids']:
+        #         track_query_target_mask[ind[1].eq(i).nonzero()[0]] = True
+
+        #     track_query_target_masks.append(track_query_target_mask)
+        # track_query_target_masks = torch.cat(track_query_target_masks)
+
+        # losses['loss_bbox_track_queries'] = loss_bbox[track_query_target_masks].sum() / num_boxes
+        # losses['loss_bbox_object_queries'] = loss_bbox[~track_query_target_masks].sum() / num_boxes
+
+        # losses['loss_giou_track_queries'] = loss_giou[track_query_target_masks].sum() / num_boxes
+        # losses['loss_giou_object_queries'] = loss_giou[~track_query_target_masks].sum() / num_boxes
+
         return losses
 
     def loss_masks(self, outputs, targets, indices, num_boxes):

src/trackformer/models/detr_tracking.py

@@ -185,8 +185,9 @@ def forward(self, samples: NestedTensor, targets: list = None, prev_features=Non
         if targets is not None and not self._tracking:
             prev_targets = [target['prev_target'] for target in targets]
 
-            if self.training: # and random.uniform(0, 1) < 0.5:
-            # if self.training:
+            # if self.training and random.uniform(0, 1) < 0.5:
+            if self.training:
+            # if True:
                 backprop_context = torch.no_grad
                 if self._backprop_prev_frame:
                     backprop_context = nullcontext

src/trackformer/models/matcher.py

@@ -19,7 +19,7 @@ class HungarianMatcher(nn.Module):
     """
 
     def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1,
-                 focal_loss: bool = False, focal_alpha: float = 0.25):
+                 focal_loss: bool = False, focal_alpha: float = 0.25, focal_gamma: float = 2.0):
         """Creates the matcher
 
         Params:
@@ -35,6 +35,7 @@ def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float
         self.cost_giou = cost_giou
         self.focal_loss = focal_loss
         self.focal_alpha = focal_alpha
+        self.focal_gamma = focal_gamma
         assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, "all costs cant be 0"
 
     @torch.no_grad()
@@ -80,9 +81,8 @@ def forward(self, outputs, targets):
 
         # Compute the classification cost.
         if self.focal_loss:
-            gamma = 2.0
-            neg_cost_class = (1 - self.focal_alpha) * (out_prob ** gamma) * (-(1 - out_prob + 1e-8).log())
-            pos_cost_class = self.focal_alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
+            neg_cost_class = (1 - self.focal_alpha) * (out_prob ** self.focal_gamma) * (-(1 - out_prob + 1e-8).log())
+            pos_cost_class = self.focal_alpha * ((1 - out_prob) ** self.focal_gamma) * (-(out_prob + 1e-8).log())
             cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids]
         else:
             # Contrary to the loss, we don't use the NLL, but approximate it in 1 - proba[target class].
@@ -136,4 +136,5 @@ def build_matcher(args):
         cost_bbox=args.set_cost_bbox,
         cost_giou=args.set_cost_giou,
         focal_loss=args.focal_loss,
-        focal_alpha=args.focal_alpha,)
+        focal_alpha=args.focal_alpha,
+        focal_gamma=args.focal_gamma,)

src/trackformer/util/misc.py

@@ -536,7 +536,7 @@ def dice_loss(inputs, targets, num_boxes):
     return loss.sum() / num_boxes
 
 
-def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2, query_mask=None):
+def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2, query_mask=None, reduction=True):
     """
     Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.
     Args:
@@ -561,6 +561,9 @@ def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: f
         alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
         loss = alpha_t * loss
 
+    if not reduction:
+        return loss
+
     if query_mask is not None:
         loss = torch.stack([l[m].mean(0) for l, m in zip(loss, query_mask)])
         return loss.sum() / num_boxes

src/train.py

@@ -263,7 +263,7 @@ def match_name_keywords(n, name_keywords):
     if args.eval_only:
         _, coco_evaluator = evaluate(
             model, criterion, postprocessors, data_loader_val, device,
-            output_dir, visualizers['val'], args)
+            output_dir, visualizers['val'], args, 0)
         if args.output_dir:
             utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")