Update core/networks.py

core/networks.py

@@ -35,44 +35,44 @@ def group_norm(features):
 
 class Backbone(nn.Module, ABC_Model):
     def __init__(self, model_name, num_classes=20, mode='fix', segmentation=False):
-    super().__init__()
-
-    self.mode = mode
-
-    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)
-
-        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)
-    else:
-        if segmentation:
-            dilation, dilated = 4, True
+        super().__init__()
+        
+        self.mode = mode
+        
+        if self.mode == 'fix':
+            self.norm_fn = FixedBatchNorm
         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)
+            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)
+        
+            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)
+        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)
 
 class Classifier(Backbone):
     def __init__(self, model_name, state_path, num_classes=20, mode='fix'):