add new bn statistics & results

Author: Tete Xiao

README.md

@@ -67,11 +67,11 @@ IMPORTANT: We use our self-trained base model on ImageNet. The model takes the i
     </tr>
     <tr>
         <td rowspan="2">ResNet-50_dilated8 + psp_bilinear_deepsup</td>
-        <td>No</td><td>40.60</td><td>79.66</td><td>60.13</td>
+        <td>No</td><td>41.26</td><td>79.73</td><td>60.50</td>
         <td rowspan="2">33.4 hours</td>
     </tr>
     <tr>
-        <td>Yes</td><td>41.31</td><td>80.14</td><td>60.73</td>
+        <td>Yes</td><td>42.04</td><td>80.23</td><td>61.14</td>
     </tr>
     <tr>
         <td>ResNet-101_dilated8 + c1_bilinear_deepsup</td>

lib/nn/modules/batchnorm.py

@@ -36,7 +36,7 @@ def _unsqueeze_ft(tensor):
 
 
 class _SynchronizedBatchNorm(_BatchNorm):
-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+    def __init__(self, num_features, eps=1e-5, momentum=0.001, affine=True):
         super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine)
 
         self._sync_master = SyncMaster(self._data_parallel_master)
@@ -45,6 +45,14 @@ def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
         self._parallel_id = None
         self._slave_pipe = None
 
+        # customed batch norm statistics
+        self._iter = 1
+        self._moving_average_fraction = 1. - momentum
+        self.register_buffer('_tmp_running_mean', torch.zeros(self.num_features))
+        self.register_buffer('_tmp_running_var', torch.ones(self.num_features))
+        self._tmp_running_mean = self.running_mean.clone()
+        self._tmp_running_var = self.running_var.clone()
+
     def forward(self, input):
         # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
         if not (self._is_parallel and self.training):
@@ -108,6 +116,10 @@ def _data_parallel_master(self, intermediates):
 
         return outputs
 
+    def _add_weighted(self, dest, delta, alpha=1, beta=1, bias=0):
+        """return *dest* by `dest := dest*alpha + delta*beta + bias`"""
+        return dest * alpha + delta * beta + bias
+
     def _compute_mean_std(self, sum_, ssum, size):
         """Compute the mean and standard-deviation with sum and square-sum. This method
         also maintains the moving average on the master device."""
@@ -117,8 +129,12 @@ def _compute_mean_std(self, sum_, ssum, size):
         unbias_var = sumvar / (size - 1)
         bias_var = sumvar / size
 
-        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
-        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+        self._tmp_running_mean = self._add_weighted(self._tmp_running_mean, mean.data, alpha=self._moving_average_fraction)
+        self._tmp_running_var = self._add_weighted(self._tmp_running_var, unbias_var.data, alpha=self._moving_average_fraction)
+        self._iter = self._add_weighted(self._iter, 1, alpha=self._moving_average_fraction)
+
+        self.running_mean = self._tmp_running_mean / self._iter
+        self.running_var = self._tmp_running_var / self._iter
 
         return mean, bias_var.clamp(self.eps) ** -0.5
 

models.py

@@ -105,7 +105,7 @@ def build_encoder(self, arch='resnet50_dilated8', fc_dim=512, weights=''):
         if len(weights) > 0:
             print('Loading weights for net_encoder')
             net_encoder.load_state_dict(
-                torch.load(weights, map_location=lambda storage, loc: storage))
+                torch.load(weights, map_location=lambda storage, loc: storage), strict=False)
         return net_encoder
 
     def build_decoder(self, arch='psp_bilinear_deepsup',
@@ -138,7 +138,7 @@ def build_decoder(self, arch='psp_bilinear_deepsup',
         if len(weights) > 0:
             print('Loading weights for net_decoder')
             net_decoder.load_state_dict(
-                torch.load(weights, map_location=lambda storage, loc: storage))
+                torch.load(weights, map_location=lambda storage, loc: storage), strict=False)
         return net_decoder
 
 

resnet.py

@@ -195,7 +195,7 @@ def resnet50(pretrained=False, **kwargs):
     """
     model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
     if pretrained:
-        model.load_state_dict(load_url(model_urls['resnet50']))
+        model.load_state_dict(load_url(model_urls['resnet50']), strict=False)
     return model
 
 
@@ -207,7 +207,7 @@ def resnet101(pretrained=False, **kwargs):
     """
     model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
     if pretrained:
-        model.load_state_dict(load_url(model_urls['resnet101']))
+        model.load_state_dict(load_url(model_urls['resnet101']), strict=False)
     return model
 
 # def resnet152(pretrained=False, **kwargs):

train.py

@@ -83,11 +83,14 @@ def checkpoint(nets, history, args, epoch_num):
     dict_encoder = net_encoder.state_dict()
     dict_decoder = net_decoder.state_dict()
 
+    dict_encoder_save = {k: v for k, v in dict_encoder.items() if not (k.endswith('_tmp_running_mean') or k.endswith('tmp_running_var'))}
+    dict_decoder_save = {k: v for k, v in dict_decoder.items() if not (k.endswith('_tmp_running_mean') or k.endswith('tmp_running_var'))}
+    
     torch.save(history,
                '{}/history_{}'.format(args.ckpt, suffix_latest))
-    torch.save(dict_encoder,
+    torch.save(dict_encoder_save,
                '{}/encoder_{}'.format(args.ckpt, suffix_latest))
-    torch.save(dict_decoder,
+    torch.save(dict_decoder_save,
                '{}/decoder_{}'.format(args.ckpt, suffix_latest))
 
 
@@ -174,7 +177,7 @@ def main(args):
     # Main loop
     history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
 
-    for epoch in range(1, args.num_epoch + 1):
+    for epoch in range(args.start_epoch, args.num_epoch + 1):
         train(segmentation_module, iterator_train, optimizers, history, epoch, args)
 
         # checkpointing
@@ -214,6 +217,8 @@ def main(args):
                         help='input batch size')
     parser.add_argument('--num_epoch', default=20, type=int,
                         help='epochs to train for')
+    parser.add_argument('--start_epoch', default=1, type=int,
+                        help='epoch to start training. useful if continue from a checkpoint')
     parser.add_argument('--epoch_iters', default=5000, type=int,
                         help='iterations of each epoch (irrelevant to batch size)')
     parser.add_argument('--optim', default='SGD', help='optimizer')