Use automatic downloads of pretrained weights instead of hardcoded paths

models/config.py

@@ -1,17 +1,11 @@
 import os
 
-# here (https://github.com/pytorch/vision/tree/master/torchvision/models) to find the download link of pretrained models
-
-root = '/media/b3-542/LIBRARY/ZijunDeng/PyTorch Pretrained'
-res101_path = os.path.join(root, 'ResNet', 'resnet101-5d3b4d8f.pth')
-res152_path = os.path.join(root, 'ResNet', 'resnet152-b121ed2d.pth')
-inception_v3_path = os.path.join(root, 'Inception', 'inception_v3_google-1a9a5a14.pth')
-vgg19_bn_path = os.path.join(root, 'VggNet', 'vgg19_bn-c79401a0.pth')
-vgg16_path = os.path.join(root, 'VggNet', 'vgg16-397923af.pth')
-dense201_path = os.path.join(root, 'DenseNet', 'densenet201-4c113574.pth')
+# PyTorch will automatically download pretrained weights into `os.environ['TORCH_MODEL_ZOO']`
+# using the mechanism described here: (http://pytorch.org/docs/master/model_zoo.html)
+# Download links used are also listed here: (https://github.com/pytorch/vision/tree/master/torchvision/models)
 
 '''
 vgg16 trained using caffe
 visit this (https://github.com/jcjohnson/pytorch-vgg) to download the converted vgg16
 '''
-vgg16_caffe_path = os.path.join(root, 'VggNet', 'vgg16-caffe.pth')
+vgg16_caffe_path = os.path.join(os.environ.get('TORCH_MODEL_ZOO', '.'), 'vgg16-caffe.pth')

models/duc_hdc.py

@@ -2,9 +2,6 @@
 from torch import nn
 from torchvision import models
 
-from .config import res152_path
-
-
 class _DenseUpsamplingConvModule(nn.Module):
     def __init__(self, down_factor, in_dim, num_classes):
         super(_DenseUpsamplingConvModule, self).__init__()
@@ -26,9 +23,7 @@ class ResNetDUC(nn.Module):
     # the size of image should be multiple of 8
     def __init__(self, num_classes, pretrained=True):
         super(ResNetDUC, self).__init__()
-        resnet = models.resnet152()
-        if pretrained:
-            resnet.load_state_dict(torch.load(res152_path))
+        resnet = models.resnet152(pretrained=pretrained)
         self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool)
         self.layer1 = resnet.layer1
         self.layer2 = resnet.layer2

models/fcn16s.py

@@ -9,9 +9,7 @@
 class FCN16VGG(nn.Module):
     def __init__(self, num_classes, pretrained=True):
         super(FCN16VGG, self).__init__()
-        vgg = models.vgg16()
-        if pretrained:
-            vgg.load_state_dict(torch.load(vgg16_caffe_path))
+        vgg = models.vgg16(pretrained=pretrained)
         features, classifier = list(vgg.features.children()), list(vgg.classifier.children())
 
         features[0].padding = (100, 100)

models/fcn32s.py

@@ -3,15 +3,12 @@
 from torchvision import models
 
 from ..utils import get_upsampling_weight
-from .config import vgg16_caffe_path
 
 
 class FCN32VGG(nn.Module):
     def __init__(self, num_classes, pretrained=True):
         super(FCN32VGG, self).__init__()
-        vgg = models.vgg16()
-        if pretrained:
-            vgg.load_state_dict(torch.load(vgg16_caffe_path))
+        vgg = models.vgg16(pretrained=pretrained)
         features, classifier = list(vgg.features.children()), list(vgg.classifier.children())
 
         features[0].padding = (100, 100)

models/fcn8s.py

@@ -3,20 +3,22 @@
 from torchvision import models
 
 from ..utils import get_upsampling_weight
-from .config import vgg16_path, vgg16_caffe_path
+from .config import vgg16_caffe_path
 
 
 # This is implemented in full accordance with the original one (https://github.com/shelhamer/fcn.berkeleyvision.org)
 class FCN8s(nn.Module):
     def __init__(self, num_classes, pretrained=True, caffe=False):
         super(FCN8s, self).__init__()
-        vgg = models.vgg16()
-        if pretrained:
-            if caffe:
-                # load the pretrained vgg16 used by the paper's author
-                vgg.load_state_dict(torch.load(vgg16_caffe_path))
-            else:
-                vgg.load_state_dict(torch.load(vgg16_path))
+
+        if pretrained and caffe:
+            vgg = models.vgg16()
+            # load the pretrained vgg16 used by the paper's author
+            vgg.load_state_dict(torch.load(vgg16_caffe_path))
+        else:
+            # if pretrained, load the weights from PyTorch model zoo
+            vgg = models.vgg16(pretrained=pretrained)
+
         features, classifier = list(vgg.features.children()), list(vgg.classifier.children())
 
         '''

models/gcn.py

@@ -4,7 +4,6 @@
 from torchvision import models
 
 from ..utils import initialize_weights
-from .config import res152_path
 
 
 # many are borrowed from https://github.com/ycszen/pytorch-ss/blob/master/gcn.py
@@ -52,9 +51,7 @@ class GCN(nn.Module):
     def __init__(self, num_classes, input_size, pretrained=True):
         super(GCN, self).__init__()
         self.input_size = input_size
-        resnet = models.resnet152()
-        if pretrained:
-            resnet.load_state_dict(torch.load(res152_path))
+        resnet = models.resnet152(pretrained=pretrained)
         self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu)
         self.layer1 = nn.Sequential(resnet.maxpool, resnet.layer1)
         self.layer2 = resnet.layer2

models/psp_net.py

@@ -5,7 +5,6 @@
 
 from ..utils import initialize_weights
 from ..utils.misc import Conv2dDeformable
-from .config import res101_path
 
 
 class _PyramidPoolingModule(nn.Module):
@@ -34,9 +33,7 @@ class PSPNet(nn.Module):
     def __init__(self, num_classes, pretrained=True, use_aux=True):
         super(PSPNet, self).__init__()
         self.use_aux = use_aux
-        resnet = models.resnet101()
-        if pretrained:
-            resnet.load_state_dict(torch.load(res101_path))
+        resnet = models.resnet101(pretrained=pretrained)
         self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool)
         self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
 

models/seg_net.py

@@ -3,7 +3,6 @@
 from torchvision import models
 
 from ..utils import initialize_weights
-from .config import vgg19_bn_path
 
 
 class _DecoderBlock(nn.Module):
@@ -35,9 +34,7 @@ def forward(self, x):
 class SegNet(nn.Module):
     def __init__(self, num_classes, pretrained=True):
         super(SegNet, self).__init__()
-        vgg = models.vgg19_bn()
-        if pretrained:
-            vgg.load_state_dict(torch.load(vgg19_bn_path))
+        vgg = models.vgg19_bn(pretrained=pretrained)
         features = list(vgg.features.children())
         self.enc1 = nn.Sequential(*features[0:7])
         self.enc2 = nn.Sequential(*features[7:14])