support other value of ignore label

Author: CoinCheung

lib/base_dataset.py

@@ -24,6 +24,7 @@ def __init__(self, dataroot, annpath, trans_func=None, mode='train'):
         self.mode = mode
         self.trans_func = trans_func
 
+        self.lb_ignore = -100
         self.lb_map = None
 
         with open(annpath, 'r') as fr:

lib/coco.py

@@ -48,9 +48,9 @@
 class CocoStuff(BaseDataset):
 
     def __init__(self, dataroot, annpath, trans_func=None, mode='train'):
-        super(CocoStuff, self).__init__(dataroot, annpath, trans_func, mode)
+        super(CocoStuff, self).__init__(
+                dataroot, annpath, trans_func, mode)
         self.n_cats = 171 # 91 stuff, 91 thing, 11 of thing have no annos
-        self.lb_ignore = 255
 
         ## label mapping, remove non-existing labels
         missing = [11, 25, 28, 29, 44, 65, 67, 68, 70, 82, 90]

lib/ohem_ce_loss.py

@@ -10,30 +10,30 @@
 #  import ohem_cpp
 #  class OhemCELoss(nn.Module):
 #
-#      def __init__(self, thresh, ignore_lb=255):
+#      def __init__(self, thresh, lb_ignore=255):
 #          super(OhemCELoss, self).__init__()
 #          self.score_thresh = thresh
-#          self.ignore_lb = ignore_lb
-#          self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='mean')
+#          self.lb_ignore = lb_ignore
+#          self.criteria = nn.CrossEntropyLoss(ignore_index=lb_ignore, reduction='mean')
 #
 #      def forward(self, logits, labels):
-#          n_min = labels[labels != self.ignore_lb].numel() // 16
+#          n_min = labels[labels != self.lb_ignore].numel() // 16
 #          labels = ohem_cpp.score_ohem_label(
-#                  logits, labels, self.ignore_lb, self.score_thresh, n_min).detach()
+#                  logits, labels, self.lb_ignore, self.score_thresh, n_min).detach()
 #          loss = self.criteria(logits, labels)
 #          return loss
 
 
 class OhemCELoss(nn.Module):
 
-    def __init__(self, thresh, ignore_lb=255):
+    def __init__(self, thresh, lb_ignore=255):
         super(OhemCELoss, self).__init__()
         self.thresh = -torch.log(torch.tensor(thresh, requires_grad=False, dtype=torch.float)).cuda()
-        self.ignore_lb = ignore_lb
-        self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='none')
+        self.lb_ignore = lb_ignore
+        self.criteria = nn.CrossEntropyLoss(ignore_index=lb_ignore, reduction='none')
 
     def forward(self, logits, labels):
-        n_min = labels[labels != self.ignore_lb].numel() // 16
+        n_min = labels[labels != self.lb_ignore].numel() // 16
         loss = self.criteria(logits, labels).view(-1)
         loss_hard = loss[loss > self.thresh]
         if loss_hard.numel() < n_min:

tools/train_amp.py

@@ -52,7 +52,7 @@ def parse_args():
 
 
 
-def set_model():
+def set_model(lb_ignore=255):
     logger = logging.getLogger()
     net = model_factory[cfg.model_type](cfg.n_cats)
     if not args.finetune_from is None:
@@ -61,8 +61,9 @@ def set_model():
     if cfg.use_sync_bn: net = nn.SyncBatchNorm.convert_sync_batchnorm(net)
     net.cuda()
     net.train()
-    criteria_pre = OhemCELoss(0.7)
-    criteria_aux = [OhemCELoss(0.7) for _ in range(cfg.num_aux_heads)]
+    criteria_pre = OhemCELoss(0.7, lb_ignore)
+    criteria_aux = [OhemCELoss(0.7, lb_ignore)
+            for _ in range(cfg.num_aux_heads)]
     return net, criteria_pre, criteria_aux
 
 
@@ -126,7 +127,7 @@ def train():
     dl = get_data_loader(cfg, mode='train', distributed=is_dist)
 
     ## model
-    net, criteria_pre, criteria_aux = set_model()
+    net, criteria_pre, criteria_aux = set_model(dl.dataset.lb_ignore)
 
     ## optimizer
     optim = set_optimizer(net)