code cleanup

Author: chenxi116

deeplab.py

@@ -4,7 +4,7 @@
 import torch.utils.model_zoo as model_zoo
 
 
-__all__ = ['ResNet', 'resnet50', 'resnet101', 'resnet152', 'pnasnet5', 'hdarts']
+__all__ = ['ResNet', 'resnet50', 'resnet101', 'resnet152']
 
 
 model_urls = {
@@ -115,13 +115,9 @@ def forward(self, x):
 
 class ResNet(nn.Module):
 
-    def __init__(self, block, layers, num_classes, num_groups=None, sync_bn=False, beta=False):
+    def __init__(self, block, layers, num_classes, num_groups=None, beta=False):
         self.inplanes = 64
-        if sync_bn:
-            from encoding.nn import BatchNorm2d
-        else:
-            from torch.nn import BatchNorm2d
-        self._norm = lambda planes, momentum=0.05: BatchNorm2d(planes, momentum=momentum) if num_groups is None else nn.GroupNorm(num_groups, planes)
+        self._norm = lambda planes, momentum=0.05: nn.BatchNorm2d(planes, momentum=momentum) if num_groups is None else nn.GroupNorm(num_groups, planes)
 
         super(ResNet, self).__init__()
         if not beta:
@@ -146,7 +142,7 @@ def __init__(self, block, layers, num_classes, num_groups=None, sync_bn=False, b
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                 m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, BatchNorm2d) or isinstance(m, nn.GroupNorm):
+            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.GroupNorm):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()
 

main.py

@@ -40,8 +40,6 @@
                     help='train from scratch')
 parser.add_argument('--freeze_bn', action='store_true', default=False,
                     help='freeze batch normalization parameters')
-parser.add_argument('--sync_bn', action='store_true', default=False,
-                    help='sync batch normalization across gpu')
 parser.add_argument('--beta', action='store_true', default=False,
                     help='resnet101 beta')
 parser.add_argument('--crop_size', type=int, default=513,
@@ -70,39 +68,27 @@ def main():
         pretrained=(not args.scratch),
         num_classes=len(dataset.CLASSES),
         num_groups=args.groups,
-        sync_bn=args.sync_bn,
         beta=args.beta)
   else:
     raise ValueError('Unknown backbone: {}'.format(args.backbone))
 
   if args.train:
     criterion = nn.CrossEntropyLoss(ignore_index=255)
-    if args.sync_bn:
-      from encoding.parallel import DataParallelModel, DataParallelCriterion
-      criterion = DataParallelCriterion(criterion).cuda()
-      model = DataParallelModel(model).cuda()
-    else:
-      model = nn.DataParallel(model).cuda()
+    model = nn.DataParallel(model).cuda()
     model.train()
     if args.freeze_bn:
       for m in model.modules():
         if isinstance(m, nn.BatchNorm2d):
           m.eval()
           m.weight.requires_grad = False
           m.bias.requires_grad = False
-    if args.backbone == 'resnet101':
-      backbone_params = (
-          list(model.module.conv1.parameters()) +
-          list(model.module.bn1.parameters()) +
-          list(model.module.layer1.parameters()) +
-          list(model.module.layer2.parameters()) +
-          list(model.module.layer3.parameters()) +
-          list(model.module.layer4.parameters()))
-    else:
-      backbone_params = (
-          list(model.module.stem0.parameters()) +
-          list(model.module.stem1.parameters()) +
-          list(model.module.cells.parameters()))
+    backbone_params = (
+        list(model.module.conv1.parameters()) +
+        list(model.module.bn1.parameters()) +
+        list(model.module.layer1.parameters()) +
+        list(model.module.layer2.parameters()) +
+        list(model.module.layer3.parameters()) +
+        list(model.module.layer4.parameters()))
     last_params = list(model.module.aspp.parameters())
     optimizer = optim.SGD([
       {'params': filter(lambda p: p.requires_grad, backbone_params)},