models: Implement load_backbone function

Author: synml

models/backbone.py

@@ -1,3 +1,5 @@
+import os
+
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.tensorboard
@@ -51,6 +53,16 @@ def make_initial_conv(self, in_channels: int, out_channels: int):
         )
 
 
+def load_backbone(num_classes: int, pretrained=False):
+    model = Backbone(num_classes)
+    if pretrained:
+        if os.path.exists('weights/Backbone_best.pth'):
+            model.load_state_dict(torch.load('weights/Backbone_best.pth'))
+        else:
+            print('FileNotFound: pretrained_weights (Backbone)')
+    return model
+
+
 if __name__ == '__main__':
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = Backbone(8).to(device)

models/proposed.py

@@ -4,7 +4,7 @@
 import torch.utils.tensorboard
 import torchsummary
 
-import utils
+import models.backbone
 
 
 # ASPP(Atrous Spatial Pyramid Pooling) Module
@@ -59,7 +59,7 @@ class Proposed(nn.Module):
     def __init__(self, num_classes: int):
         super(Proposed, self).__init__()
         # Backbone
-        backbone = utils.get_model('Backbone', num_classes, 'weights/Backbone_best.pth')
+        backbone = models.backbone.load_backbone(num_classes, pretrained=True)
         self.initial_conv = backbone.initial_conv
         self.encode1 = backbone.layer1  # 64
         self.encode2 = backbone.layer2  # 128, 1/2