merge pytorch0.3.1

Author: pangwong

demo.sh

@@ -11,7 +11,8 @@
 #
 # 4. uncomment to train
 #
-# python multi_label_classifier.py --dir "./test/celeba/" --mode "Train" --name "test" --batch_size 64 --gpu_ids 0 --input_channel 3 --load_size 144 --input_size 128 --mean [0,0,0] --std [1,1,1] --ratio "[0.94, 0.03, 0.03]" --shuffle --load_thread 8 --sum_epoch 20 --lr_decay_in_epoch 4 --display_port 8900 --validate_ratio 0.1 --top_k "(1,)" --score_thres 0.1 --display_train_freq 20 --display_validate_freq 20 --save_epoch_freq 1  --display_image_ratio 0.2
+# python multi_label_classifier.py --dir "./test/celeba/" --mode "Train" --model VGG16 --name "test" --batch_size 64 --gpu_ids 0 --input_channel 3 --load_size 144 --input_size 128 --mean [0,0,0] --std [1,1,1] --ratio "[0.94, 0.03, 0.03]" --shuffle --load_thread 8 --sum_epoch 20 --lr_decay_in_epoch 4 --display_port 8900 --validate_ratio 0.1 --top_k "(1,)" --score_thres 0.1 --display_train_freq 20 --display_validate_freq 20 --save_epoch_freq 1  --display_image_ratio 0.2
+#
 # 5. open localhost:8900 on your browser and you will see loss and accuracy curves and training images samples later on.
 #
 #--------------test--------------

models/alexnet.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 import torch.utils.model_zoo as model_zoo
-
+from build_model import LoadPretrainedModel
 
 __all__ = ['AlexNet', 'alexnet']
 
@@ -89,3 +89,15 @@ def alexnet(pretrained=False, **kwargs):
     if pretrained:
         model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
     return model
+
+def AlexnetTemplet(input_channel, pretrained=False, **kwargs):
+    r"""AlexNet model architecture from the
+    `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = AlexNetTemplet(input_channel)
+    if pretrained:
+        model_dict = LoadPretrainedModel(model, model_zoo.load_url(model_urls['alexnet']))
+        model.load_state_dict(model_dict)
+    return model

models/build_model.py

@@ -18,6 +18,11 @@ def forward(self, x):
             outs.append(out)
         return outs
 
+def LoadPretrainedModel(model, pretrained_state_dict):
+    model_dict = model.state_dict()
+    union_dict = {k : v for k,v in pretrained_state_dict.iteritems() if k in model_dict}
+    model_dict.update(union_dict)
+    return model_dict
 
 def BuildMultiLabelModel(basemodel, basemodel_output, num_classes):
     return MultiLabelModel(basemodel, basemodel_output, num_classes)

models/resnet.py

@@ -1,6 +1,7 @@
 import torch.nn as nn
 import math
 import torch.utils.model_zoo as model_zoo
+from build_model import LoadPretrainedModel
 
 
 __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
@@ -229,7 +230,8 @@ def Resnet18Templet(input_channel, pretrained=False, **kwargs):
     """
     model = ResNetTemplet(BasicBlock, [2, 2, 2, 2], input_channel, **kwargs)
     if pretrained:
-        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+        model_dict = LoadPretrainedModel(model, model_zoo.load_url(model_urls['resnet18']))
+        model.load_state_dict(model_dict)
     return model
 
 
@@ -253,7 +255,8 @@ def Resnet34Templet(input_channel, pretrained=False, **kwargs):
     """
     model = ResNetTemplet(BasicBlock, [3, 4, 6, 3], input_channel, **kwargs)
     if pretrained:
-        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+        model_dict = LoadPretrainedModel(model, model_zoo.load_url(model_urls['resnet34']))
+        model.load_state_dict(model_dict)
     return model
 
 
@@ -277,7 +280,8 @@ def Resnet50Templet(input_channel, pretrained=False, **kwargs):
     """
     model = ResNetTemplet(Bottleneck, [3, 4, 6, 3], input_channel, **kwargs)
     if pretrained:
-        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+        model_dict = LoadPretrainedModel(model, model_zoo.load_url(model_urls['resnet50']))
+        model.load_state_dict(model_dict)
     return model
 
 
@@ -300,7 +304,8 @@ def Resnet101Templet(input_channel, pretrained=False, **kwargs):
     """
     model = ResNetTemplet(Bottleneck, [3, 4, 23, 3], input_channel, **kwargs)
     if pretrained:
-        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+        model_dict = LoadPretrainedModel(model, model_zoo.load_url(model_urls['resnet101']))
+        model.load_state_dict(model_dict)
     return model
 
 def resnet152(pretrained=False, **kwargs):
@@ -322,6 +327,7 @@ def Resnet152Templet(input_channel, pretrained=False, **kwargs):
     """
     model = ResNetTemplet(Bottleneck, [3, 8, 36, 3], input_channel, **kwargs)
     if pretrained:
-        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+        model_dict = LoadPretrainedModel(model, model_zoo.load_url(model_urls['resnet152']))
+        model.load_state_dict(model_dict)
     return model