compute loss inside model

Author: lufficc

ssd/engine/trainer.py

@@ -47,7 +47,6 @@ def do_train(cfg, model,
              data_loader,
              optimizer,
              scheduler,
-             criterion,
              device,
              args):
     logger = logging.getLogger("SSD.trainer")
@@ -74,14 +73,13 @@ def do_train(cfg, model,
         labels = labels.to(device)
 
         optimizer.zero_grad()
-        confidence, locations = model(images)
-        regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
+        loss_dict = model(images, targets=(boxes, labels))
 
         # reduce losses over all GPUs for logging purposes
-        loss_dict_reduced = reduce_loss_dict({'regression_loss': regression_loss, 'classification_loss': classification_loss})
+        loss_dict_reduced = reduce_loss_dict(loss_dict)
         losses_reduced = sum(loss for loss in loss_dict_reduced.values())
 
-        loss = regression_loss + classification_loss
+        loss = sum(loss for loss in loss_dict.values())
         loss.backward()
         optimizer.step()
         trained_time += time.time() - end

ssd/modeling/ssd.py

@@ -3,6 +3,7 @@
 import numpy as np
 import torch.nn.functional as F
 
+from ssd.modeling.multibox_loss import MultiBoxLoss
 from ssd.module import L2Norm
 from ssd.module.prior_box import PriorBox
 from ssd.utils import box_utils
@@ -24,6 +25,7 @@ def __init__(self, cfg,
         self.classification_headers = classification_headers
         self.regression_headers = regression_headers
         self.l2_norm = L2Norm(512, scale=20)
+        self.criterion = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
         self.priors = None
         self.reset_parameters()
 
@@ -38,7 +40,7 @@ def weights_init(m):
         self.classification_headers.apply(weights_init)
         self.regression_headers.apply(weights_init)
 
-    def forward(self, x):
+    def forward(self, x, targets=None):
         sources = []
         confidences = []
         locations = []
@@ -68,17 +70,24 @@ def forward(self, x):
         locations = locations.view(locations.size(0), -1, 4)
 
         if not self.training:
+            # when evaluating, decode predictions
             if self.priors is None:
                 self.priors = PriorBox(self.cfg)().to(locations.device)
             confidences = F.softmax(confidences, dim=2)
             boxes = box_utils.convert_locations_to_boxes(
                 locations, self.priors, self.cfg.MODEL.CENTER_VARIANCE, self.cfg.MODEL.SIZE_VARIANCE
             )
             boxes = box_utils.center_form_to_corner_form(boxes)
-
             return confidences, boxes
         else:
-            return confidences, locations
+            # when training, compute losses
+            gt_boxes, gt_labels = targets
+            regression_loss, classification_loss = self.criterion(confidences, locations, gt_labels, gt_boxes)
+            loss_dict = dict(
+                regression_loss=regression_loss,
+                classification_loss=classification_loss,
+            )
+            return loss_dict
 
     def init_from_base_net(self, model):
         vgg_weights = torch.load(model, map_location=lambda storage, loc: storage)

train_ssd.py

@@ -12,7 +12,6 @@
 from ssd.engine.inference import do_evaluation
 from ssd.engine.trainer import do_train
 from ssd.modeling.data_preprocessing import TrainAugmentation
-from ssd.modeling.multibox_loss import MultiBoxLoss
 from ssd.modeling.ssd import MatchPrior
 from ssd.modeling.vgg_ssd import build_ssd_model
 from ssd.module.prior_box import PriorBox
@@ -43,10 +42,6 @@ def train(cfg, args):
     # -----------------------------------------------------------------------------
     lr = cfg.SOLVER.LR * args.num_gpus  # scale by num gpus
     optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY)
-    # -----------------------------------------------------------------------------
-    # Criterion
-    # -----------------------------------------------------------------------------
-    criterion = MultiBoxLoss(neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
 
     # -----------------------------------------------------------------------------
     # Scheduler
@@ -73,7 +68,7 @@ def train(cfg, args):
     batch_sampler = samplers.IterationBasedBatchSampler(batch_sampler, num_iterations=cfg.SOLVER.MAX_ITER // args.num_gpus)
     train_loader = DataLoader(train_dataset, num_workers=4, batch_sampler=batch_sampler, pin_memory=True)
 
-    return do_train(cfg, model, train_loader, optimizer, scheduler, criterion, device, args)
+    return do_train(cfg, model, train_loader, optimizer, scheduler, device, args)
 
 
 def main():