utils: Add parameter of get model()

Author: synml

train.py

@@ -19,7 +19,7 @@
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.get_model(config).to(device)
+    model = utils.get_model(config, pretrained_backbone=True).to(device)
 
     # 3. Loss function, optimizer, lr scheduler, scaler
     criterion = nn.CrossEntropyLoss()

utils.py

@@ -24,7 +24,7 @@ def load_config():
     return config
 
 
-def get_model(config: dict, pretrained=False) -> torch.nn.Module:
+def get_model(config: dict, pretrained=False, pretrained_backbone=False) -> torch.nn.Module:
     assert isinstance(pretrained, bool)
     assert config['dataset']['num_classes'] == 20 or config['dataset']['num_classes'] == 8
 
@@ -33,7 +33,10 @@ 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'], config['Backbone']['pretrained_weights'])
+        if pretrained_backbone:
+            model = models.proposed.Proposed(config['dataset']['num_classes'], config['Backbone']['pretrained_weights'])
+        else:
+            model = models.proposed.Proposed(config['dataset']['num_classes'])
     else:
         raise NameError('Wrong model name.')