models: Remove hard coding backbone weights path

Author: synml

models/backbone.py

@@ -47,11 +47,11 @@ def make_initial_conv(self, in_channels: int, out_channels: int):
         )
 
 
-def load_backbone(num_classes: int, pretrained=False):
+def load_backbone(num_classes: int, pretrained_weights: str = None):
     model = Backbone(num_classes)
-    if pretrained:
-        if os.path.exists('weights/Backbone_val_best.pth'):
-            model.load_state_dict(torch.load('weights/Backbone_val_best.pth'))
+    if pretrained_weights is not None:
+        if os.path.exists(pretrained_weights):
+            model.load_state_dict(torch.load(pretrained_weights))
         else:
             print('FileNotFound: pretrained_weights (Backbone)')
     return model

models/proposed.py

@@ -8,10 +8,10 @@
 
 
 class Proposed(nn.Module):
-    def __init__(self, num_classes: int):
+    def __init__(self, num_classes: int, backbone_pretrained_weights: str = None):
         super(Proposed, self).__init__()
         # Backbone
-        backbone = models.backbone.load_backbone(num_classes, pretrained=True)
+        backbone = models.backbone.load_backbone(num_classes, backbone_pretrained_weights)
         self.initial_conv = backbone.initial_conv
         self.encode1 = backbone.layer1  # 64
         self.encode2 = backbone.layer2  # 128, 1/2

utils.py

@@ -33,7 +33,7 @@ def get_model(config: dict, pretrained=False) -> torch.nn.Module:
     elif config['model'] == 'Backbone':
         model = models.backbone.Backbone(config['dataset']['num_classes'])
     elif config['model'] == 'Proposed':
-        model = models.proposed.Proposed(config['dataset']['num_classes'])
+        model = models.proposed.Proposed(config['dataset']['num_classes'], config['Backbone']['pretrained_weights'])
     else:
         raise NameError('Wrong model name.')