Move init_module in Network class to share the initialization.

Author: ruotianluo

lib/nets/network.py

@@ -310,6 +310,21 @@ def create_architecture(self, num_classes, tag=None,
     # Initialize layers
     self._init_modules()
 
+  def _init_modules(self):
+    self._init_head_tail()
+
+    # rpn
+    self.rpn_net = nn.Conv2d(self._net_conv_channels, cfg.RPN_CHANNELS, [3, 3], padding=1)
+
+    self.rpn_cls_score_net = nn.Conv2d(cfg.RPN_CHANNELS, self._num_anchors * 2, [1, 1])
+    
+    self.rpn_bbox_pred_net = nn.Conv2d(cfg.RPN_CHANNELS, self._num_anchors * 4, [1, 1])
+
+    self.cls_score_net = nn.Linear(self._fc7_channels, self._num_classes)
+    self.bbox_pred_net = nn.Linear(self._fc7_channels, self._num_classes * 4)
+
+    self.init_weights()
+
   def _run_summary_op(self, val=False):
     """
     Run the summary operator: feed the placeholders with corresponding newtork outputs(activations)

lib/nets/resnet_v1.py

@@ -211,6 +211,8 @@ def __init__(self, num_layers=50):
     self._feat_stride = [16, ]
     self._feat_compress = [1. / float(self._feat_stride[0]), ]
     self._num_layers = num_layers
+    self._net_conv_channels = 1024
+    self._fc7_channels = 2048
 
   def _crop_pool_layer(self, bottom, rois):
     return Network._crop_pool_layer(self, bottom, rois, cfg.RESNET.MAX_POOL)
@@ -225,7 +227,7 @@ def _head_to_tail(self, pool5):
     fc7 = self.resnet.layer4(pool5).mean(3).mean(2) # average pooling after layer4
     return fc7
 
-  def _init_modules(self):
+  def _init_head_tail(self):
     # choose different blocks for different number of layers
     if self._num_layers == 50:
       self.resnet = resnet50()
@@ -262,18 +264,6 @@ def set_bn_fix(m):
     self._layers['head'] = nn.Sequential(self.resnet.conv1, self.resnet.bn1,self.resnet.relu, 
       self.resnet.maxpool,self.resnet.layer1,self.resnet.layer2,self.resnet.layer3)
 
-    # rpn
-    self.rpn_net = nn.Conv2d(1024, 512, [3, 3], padding=1)
-
-    self.rpn_cls_score_net = nn.Conv2d(512, self._num_anchors * 2, [1, 1])
-    
-    self.rpn_bbox_pred_net = nn.Conv2d(512, self._num_anchors * 4, [1, 1])
-
-    self.cls_score_net = nn.Linear(2048, self._num_classes)
-    self.bbox_pred_net = nn.Linear(2048, self._num_classes * 4)
-
-    self.init_weights()
-
   def train(self, mode=True):
     # Override train so that the training mode is set as we want
     nn.Module.train(self, mode)

lib/nets/vgg16.py

@@ -22,8 +22,10 @@ def __init__(self):
     Network.__init__(self)
     self._feat_stride = [16, ]
     self._feat_compress = [1. / float(self._feat_stride[0]), ]
+    self._net_conv_channels = 512
+    self._fc7_channels = 4096
 
-  def _init_modules(self):
+  def _init_head_tail(self):
     self.vgg = models.vgg16()
     # Remove fc8
     self.vgg.classifier = nn.Sequential(*list(self.vgg.classifier._modules.values())[:-1])
@@ -35,18 +37,6 @@ def _init_modules(self):
     # not using the last maxpool layer
     self._layers['head'] = nn.Sequential(*list(self.vgg.features._modules.values())[:-1])
 
-    # rpn
-    self.rpn_net = nn.Conv2d(512, 512, [3, 3], padding=1)
-
-    self.rpn_cls_score_net = nn.Conv2d(512, self._num_anchors * 2, [1, 1])
-    
-    self.rpn_bbox_pred_net = nn.Conv2d(512, self._num_anchors * 4, [1, 1])
-
-    self.cls_score_net = nn.Linear(4096, self._num_classes)
-    self.bbox_pred_net = nn.Linear(4096, self._num_classes * 4)
-
-    self.init_weights()
-
   def _image_to_head(self):
     net_conv = self._layers['head'](self._image)
     self._act_summaries['conv'] = net_conv