add freeze_parameter for model components

jwyang · jwyang · commit e23cc54f8c45 · 2019-08-17T17:05:17.000-04:00
diff --git a/configs/baseline_res101.yaml b/configs/baseline_res101.yaml
@@ -1,19 +1,24 @@
 DATASET:
   NAME: "vg"
   MODE: "benchmark"
-  TRAIN_BATCH_SIZE: 6
+  TRAIN_BATCH_SIZE: 16
   TEST_BATCH_SIZE: 1
 MODEL:
-  WEIGHT: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_IMG: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_DET: "checkpoints/vg_benchmark_object/R-101-C4/faster_rcnn/BatchSize_6/Base_LR_0.005/checkpoint_0099999.pth"
+  RELATION_ON: True
   ALGORITHM: "sg_baseline"
   USE_FREQ_PRIOR: False
   BACKBONE:
     CONV_BODY: "R-101-C4"
+    FREEZE_PARAMETER: True
+  RPN:
+    FREEZE_PARAMETER: True
   ROI_HEADS:
-    BATCH_SIZE_PER_IMAGE: 384
+    BATCH_SIZE_PER_IMAGE: 512
   ROI_BOX_HEAD:
     NUM_CLASSES: 151
-  RELATION_ON: True
+    FREEZE_PARAMETER: True
   ROI_RELATION_HEAD:
     NUM_CLASSES: 51
 SOLVER:
diff --git a/configs/faster_rcnn_res101.yaml b/configs/faster_rcnn_res101.yaml
@@ -5,7 +5,7 @@ DATASET:
   TRAIN_BATCH_SIZE: 6
   TEST_BATCH_SIZE: 1
 MODEL:
-  WEIGHT: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_IMG: "catalog://ImageNetPretrained/MSRA/R-101"
   ALGORITHM: "faster_rcnn"
   BACKBONE:
     CONV_BODY: "R-101-C4"
diff --git a/configs/grcnn_res101.yaml b/configs/grcnn_res101.yaml
@@ -0,0 +1,26 @@
+DATASET:
+  NAME: "vg"
+  MODE: "benchmark"
+  TRAIN_BATCH_SIZE: 6
+  TEST_BATCH_SIZE: 1
+MODEL:
+  WEIGHT_IMG: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_DET: ""
+  ALGORITHM: "sg_grcnn"
+  BACKBONE:
+    CONV_BODY: "R-101-C4"
+    FREEZE_PARAMETER: False
+  RPN:
+    FREEZE_PARAMETER: False
+  ROI_HEADS:
+    BATCH_SIZE_PER_IMAGE: 384
+  ROI_BOX_HEAD:
+    NUM_CLASSES: 151
+    FREEZE_PARAMETER: False
+  ROI_RELATION_HEAD:
+    NUM_CLASSES: 51
+SOLVER:
+  BASE_LR: 5e-3
+  MAX_ITER: 100000
+  STEPS: (70000,90000)
+  CHECKPOINT_PERIOD: 5000
diff --git a/configs/imp_res101.yaml b/configs/imp_res101.yaml
@@ -4,15 +4,19 @@ DATASET:
   TRAIN_BATCH_SIZE: 6
   TEST_BATCH_SIZE: 1
 MODEL:
-  WEIGHT: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_IMG: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_DET: ""
   ALGORITHM: "sg_imp"
   BACKBONE:
     CONV_BODY: "R-101-C4"
+    FREEZE_PARAMETER: False
+  RPN:
+    FREEZE_PARAMETER: False
   ROI_HEADS:
     BATCH_SIZE_PER_IMAGE: 384
   ROI_BOX_HEAD:
     NUM_CLASSES: 151
-  RELATION_ON: True
+    FREEZE_PARAMETER: False
   ROI_RELATION_HEAD:
     NUM_CLASSES: 51
 SOLVER:
diff --git a/configs/motifnet_res101.yaml b/configs/motifnet_res101.yaml
@@ -0,0 +1,26 @@
+DATASET:
+  NAME: "vg"
+  MODE: "benchmark"
+  TRAIN_BATCH_SIZE: 6
+  TEST_BATCH_SIZE: 1
+MODEL:
+  WEIGHT_IMG: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_DET: ""
+  ALGORITHM: "sg_motifnet"
+  BACKBONE:
+    CONV_BODY: "R-101-C4"
+    FREEZE_PARAMETER: False
+  RPN:
+    FREEZE_PARAMETER: False
+  ROI_HEADS:
+    BATCH_SIZE_PER_IMAGE: 384
+  ROI_BOX_HEAD:
+    NUM_CLASSES: 151
+    FREEZE_PARAMETER: False
+  ROI_RELATION_HEAD:
+    NUM_CLASSES: 51
+SOLVER:
+  BASE_LR: 5e-3
+  MAX_ITER: 100000
+  STEPS: (70000,90000)
+  CHECKPOINT_PERIOD: 5000
diff --git a/configs/msdn_res101.yaml b/configs/msdn_res101.yaml
@@ -0,0 +1,26 @@
+DATASET:
+  NAME: "vg"
+  MODE: "benchmark"
+  TRAIN_BATCH_SIZE: 6
+  TEST_BATCH_SIZE: 1
+MODEL:
+  WEIGHT_IMG: "catalog://ImageNetPretrained/MSRA/R-101"
+  WEIGHT_DET: "checkpoints/vg_benchmark_object/R-101-C4/faster_rcnn/BatchSize_6/Base_LR_0.005/checkpoint_0099999.pth"
+  ALGORITHM: "sg_msdn"
+  BACKBONE:
+    CONV_BODY: "R-101-C4"
+    FREEZE_PARAMETER: False
+  RPN:
+    FREEZE_PARAMETER: False
+  ROI_HEADS:
+    BATCH_SIZE_PER_IMAGE: 384
+  ROI_BOX_HEAD:
+    NUM_CLASSES: 151
+    FREEZE_PARAMETER: False
+  ROI_RELATION_HEAD:
+    NUM_CLASSES: 51
+SOLVER:
+  BASE_LR: 5e-3
+  MAX_ITER: 100000
+  STEPS: (70000,90000)
+  CHECKPOINT_PERIOD: 5000
diff --git a/lib/config/defaults.py b/lib/config/defaults.py
@@ -46,14 +46,18 @@
 _C.MODEL.DEVICE = "cuda"
 _C.MODEL.META_ARCHITECTURE = "GeneralizedRCNN"
 _C.MODEL.CLS_AGNOSTIC_BBOX_REG = False
-_C.MODEL.WEIGHT = ""
+_C.MODEL.WEIGHT_IMG = ""    # weight loading path for imagenet pretrained model
+_C.MODEL.WEIGHT_DET = ""    # weight loading path for detector pre-trained model
+_C.MODEL.WEIGHT_SGG = ""    # weight loading path for scene graph generator pre-trained model
+
 # If the WEIGHT starts with a catalog://, like :R-50, the code will look for
 # the path in paths_catalog. Else, it will use it as the specified absolute
 # path
 _C.MODEL.BACKBONE = CN()  # Backbone options
 _C.MODEL.BACKBONE.CONV_BODY = "R-50-C4" # The backbone conv body to use # (e.g., 'FPN.add_fpn_ResNet101_conv5_body' to specify a ResNet-101-FPN backbone)
 _C.MODEL.BACKBONE.FREEZE_CONV_BODY_AT = 2 # Add StopGrad at a specified stage so the bottom layers are frozen
 _C.MODEL.BACKBONE.OUT_CHANNELS = 256 * 4
+_C.MODEL.BACKBONE.FREEZE_PARAMETER = False
 
 _C.MODEL.FPN = CN() # FPN options
 _C.MODEL.FPN.USE_GN = False
@@ -110,6 +114,8 @@
 # Custom rpn head, empty to use default conv or separable conv
 _C.MODEL.RPN.RPN_HEAD = "SingleConvRPNHead"
 
+_C.MODEL.RPN.FREEZE_PARAMETER = False
+
 # ---------------------------------------------------------------------------- #
 # ROI HEADS options
 # ---------------------------------------------------------------------------- #
@@ -160,6 +166,7 @@
 _C.MODEL.ROI_BOX_HEAD.DILATION = 1
 _C.MODEL.ROI_BOX_HEAD.CONV_HEAD_DIM = 256
 _C.MODEL.ROI_BOX_HEAD.NUM_STACKED_CONVS = 4
+_C.MODEL.ROI_BOX_HEAD.FREEZE_PARAMETER = False
 
 ''''''
 _C.MODEL.ROI_RELATION_HEAD = CN()
diff --git a/lib/scene_parser/parser.py b/lib/scene_parser/parser.py
@@ -14,7 +14,7 @@
 from .rcnn.utils.comm import synchronize, get_rank
 from .rcnn.modeling.relation_heads.relation_heads import build_roi_relation_head
 
-SCENE_PAESER_DICT = {"sg_baseline", "sg_imp"} #, "msdn": MSDN}
+SCENE_PAESER_DICT = {"sg_baseline", "sg_imp", "sg_msdn"} #, "msdn": MSDN}
 
 class SceneParser(GeneralizedRCNN):
     "Scene Parser"
@@ -25,6 +25,41 @@ def __init__(self, cfg):
         self.rel_heads = None
         if cfg.MODEL.RELATION_ON and self.cfg.MODEL.ALGORITHM in SCENE_PAESER_DICT:
             self.rel_heads = build_roi_relation_head(cfg, self.backbone.out_channels)
+        self._freeze_components(self.cfg)
+
+    def _freeze_components(self, cfg):
+        if cfg.MODEL.BACKBONE.FREEZE_PARAMETER:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+        if cfg.MODEL.RPN.FREEZE_PARAMETER:
+            for param in self.rpn.parameters():
+                param.requires_grad = False
+
+        if cfg.MODEL.ROI_BOX_HEAD.FREEZE_PARAMETER:
+            for param in self.roi_heads.box.parameters():
+                param.requires_grad = False
+
+    def train(self):
+        if self.cfg.MODEL.BACKBONE.FREEZE_PARAMETER:
+            self.backbone.eval()
+        else:
+            self.backbone.train()
+
+        if self.cfg.MODEL.RPN.FREEZE_PARAMETER:
+            self.rpn.eval()
+        else:
+            self.rpn.train()
+
+        if self.cfg.MODEL.ROI_BOX_HEAD.FREEZE_PARAMETER:
+            self.roi_heads.eval()
+        else:
+            self.roi_heads.train()
+
+        self.rel_heads.train()
+
+    def eval(self):
+        self.eval()
 
     def forward(self, images, targets=None):
         """
@@ -44,10 +79,11 @@ def forward(self, images, targets=None):
         images = to_image_list(images)
         features = self.backbone(images.tensors)
         proposals, proposal_losses = self.rpn(images, features, targets)
-
+        scene_parser_losses = {}
         if self.roi_heads:
-            x, detections, scene_parser_losses = self.roi_heads(features, proposals, targets)
+            x, detections, roi_heads_loss = self.roi_heads(features, proposals, targets)
             result = detections
+            scene_parser_losses.update(roi_heads_loss)
 
             if self.rel_heads:
                 relation_features = features
@@ -60,8 +96,8 @@ def forward(self, images, targets=None):
                     relation_features = x
                 # During training, self.box() will return the unaltered proposals as "detections"
                 # this makes the API consistent during training and testing
-                x_pairs, detection_pairs, loss_relation = self.rel_heads(relation_features, detections, targets)
-                losses.update(loss_relation)
+                x_pairs, detection_pairs, rel_heads_loss = self.rel_heads(relation_features, detections, targets)
+                scene_parser_losses.update(rel_heads_loss)
 
                 x = (x, x_pairs)
                 result = (detections, detection_pairs)
@@ -109,5 +145,6 @@ def build_scene_parser_optimizer(cfg, model, local_rank=0, distributed=False):
     save_to_disk = get_rank() == 0
     checkpointer = SceneParserCheckpointer(cfg, model, optimizer, scheduler, save_dir, save_to_disk,
         logger=logging.getLogger("scene_graph_generation.checkpointer"))
-    extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, resume=cfg.resume)
+    model_weight =cfg.MODEL.WEIGHT_DET if cfg.MODEL.WEIGHT_DET != "" else cfg.MODEL.WEIGHT_IMG
+    extra_checkpoint_data = checkpointer.load(model_weight, resume=cfg.resume)
     return optimizer, scheduler, checkpointer, extra_checkpoint_data
diff --git a/lib/scene_parser/rcnn/modeling/roi_heads/roi_heads.py b/lib/scene_parser/rcnn/modeling/roi_heads/roi_heads.py
@@ -19,11 +19,9 @@ def __init__(self, cfg, heads):
             self.keypoint.feature_extractor = self.box.feature_extractor
 
     def forward(self, features, proposals, targets=None):
-        losses = {}
         # TODO rename x to roi_box_features, if it doesn't increase memory consumption
         x, detections, loss_box = self.box(features, proposals, targets)
-        losses.update(loss_box)
-        return x, detections, losses
+        return x, detections, loss_box
 
 
 def build_roi_heads(cfg, in_channels):
diff --git a/lib/scene_parser/rcnn/utils/checkpoint.py b/lib/scene_parser/rcnn/utils/checkpoint.py
@@ -50,7 +50,7 @@ def save(self, name, **kwargs):
         self.tag_last_checkpoint(save_file)
 
     def load(self, f=None, resume=0, use_latest=True):
-        if self.has_last_checkpoint() and use_latest and resume == 0:
+        if self.has_last_checkpoint() and use_latest and resume > 0:
             # override argument with existing checkpoint
             f = self.get_last_checkpoint_file()
         elif self.has_checkpoint(resume) and resume > 0:
@@ -62,14 +62,14 @@ def load(self, f=None, resume=0, use_latest=True):
         self.logger.info("Loading checkpoint from {}".format(f))
         checkpoint = self._load_file(f)
         self._load_model(checkpoint)
-        if "optimizer" in checkpoint and self.optimizer:
+        if "optimizer" in checkpoint and self.optimizer and "sg" in f:
             self.logger.info("Loading optimizer from {}".format(f))
             self.optimizer.load_state_dict(checkpoint.pop("optimizer"))
-        if "scheduler" in checkpoint and self.scheduler:
+        if "scheduler" in checkpoint and self.scheduler and "sg" in f:
             self.logger.info("Loading scheduler from {}".format(f))
             self.scheduler.load_state_dict(checkpoint.pop("scheduler"))
+        checkpoint['iteration'] = resume  # if we load detector, the we should not use its start iteration
 
-        # return any further checkpoint data
         return checkpoint
 
     def has_last_checkpoint(self):
