Update core/networks.py

core/networks.py

@@ -34,45 +34,45 @@ def group_norm(features):
 #######################################################################
 
 class Backbone(nn.Module, ABC_Model):
-        def __init__(self, model_name, num_classes=20, mode='fix', segmentation=False):
-        super().__init__()
+    def __init__(self, model_name, num_classes=20, mode='fix', segmentation=False):
+    super().__init__()
 
-        self.mode = mode
+    self.mode = mode
 
-        if self.mode == 'fix':
-            self.norm_fn = FixedBatchNorm
-        else:
-            self.norm_fn = nn.BatchNorm2d
+    if self.mode == 'fix':
+        self.norm_fn = FixedBatchNorm
+    else:
+        self.norm_fn = nn.BatchNorm2d
 
-        if 'resnet' in model_name:
-            self.model = resnet.ResNet(resnet.Bottleneck, resnet.layers_dic[model_name], strides=(2, 2, 2, 1),
-                                       batch_norm_fn=self.norm_fn)
+    if 'resnet' in model_name:
+        self.model = resnet.ResNet(resnet.Bottleneck, resnet.layers_dic[model_name], strides=(2, 2, 2, 1),
+                                   batch_norm_fn=self.norm_fn)
 
-            state_dict = model_zoo.load_url(resnet.urls_dic[model_name])
-            state_dict.pop('fc.weight')
-            state_dict.pop('fc.bias')
+        state_dict = model_zoo.load_url(resnet.urls_dic[model_name])
+        state_dict.pop('fc.weight')
+        state_dict.pop('fc.bias')
 
-            self.model.load_state_dict(state_dict)
+        self.model.load_state_dict(state_dict)
+    else:
+        if segmentation:
+            dilation, dilated = 4, True
         else:
-            if segmentation:
-                dilation, dilated = 4, True
-            else:
-                dilation, dilated = 2, False
-
-            self.model = eval("resnest." + model_name)(pretrained=True, dilated=dilated, dilation=dilation,
-                                                       norm_layer=self.norm_fn)
-
-            del self.model.avgpool
-            del self.model.fc
-
-        self.stage1 = nn.Sequential(self.model.conv1,
-                                    self.model.bn1,
-                                    self.model.relu,
-                                    self.model.maxpool)
-        self.stage2 = nn.Sequential(self.model.layer1)
-        self.stage3 = nn.Sequential(self.model.layer2)
-        self.stage4 = nn.Sequential(self.model.layer3)
-        self.stage5 = nn.Sequential(self.model.layer4)
+            dilation, dilated = 2, False
+
+        self.model = eval("resnest." + model_name)(pretrained=True, dilated=dilated, dilation=dilation,
+                                                   norm_layer=self.norm_fn)
+
+        del self.model.avgpool
+        del self.model.fc
+
+    self.stage1 = nn.Sequential(self.model.conv1,
+                                self.model.bn1,
+                                self.model.relu,
+                                self.model.maxpool)
+    self.stage2 = nn.Sequential(self.model.layer1)
+    self.stage3 = nn.Sequential(self.model.layer2)
+    self.stage4 = nn.Sequential(self.model.layer3)
+    self.stage5 = nn.Sequential(self.model.layer4)
 
 class Classifier(Backbone):
     def __init__(self, model_name, state_path, num_classes=20, mode='fix'):