Issue#3

model/embedding.py

@@ -1,14 +1,29 @@
 import torch.nn as nn
 import torch.nn.functional as F
+from model.resnet import resnet18
+from model.resnet import resnet34
 from model.resnet import resnet50
+from model.resnet import resnet101
+from model.resnet import resnet152
 from config.base_config import cfg
 
 
 class EmbeddingResnet(nn.Module):
-    def __init__(self):
+    def __init__(self,model_name = 'resnet50'):
         super(EmbeddingResnet, self).__init__()
+        if model_name == 'resnet18':
+            resnet = resnet18(pretrained=True)
+        elif model_name == 'resnet34':
+            resnet = resnet34(pretrained=True)
+        elif model_name == 'resnet50':
+            resnet = resnet50(pretrained=True)
+        elif model_name == 'resnet101':
+            resnet = resnet101(pretrained=True)
+        elif model_name == 'resnet152':
+            resnet = resnet152(pretrained=True)
+        else:
+            raise ValueError('The model doesnt exist! Please choose from resnet18,resnet34,resnet50,resnet101,resnet152.')
 
-        resnet = resnet50(pretrained=True)
         self.features = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1,
                                       resnet.layer2, resnet.layer3, resnet.layer4, resnet.avgpool)
         # Fix blocks

model/net.py

@@ -21,7 +21,7 @@ def get_model(args, device):
     # Model
     embeddingNet = None
     if (args.dataset == 'custom') or (args.dataset == 'vggface2'):
-        embeddingNet = embedding.EmbeddingResnet()
+        embeddingNet = embedding.EmbeddingResnet(args.arch)
     elif (args.dataset == 'mnist') or (args.dataset == 'fmnist'):
         embeddingNet = embedding.EmbeddingLeNet()
     else:

train.py

@@ -158,6 +158,8 @@ def save_checkpoint(state, file_name):
 
     parser.add_argument('--train_log_step', type=int, default=100, metavar='M',
                         help='Number of iterations after which to log the loss')
+    parser.add_argument('--arch', type = str, default ='resnet50',
+                       help="model architecture . default resnet50 . Choose from 'resnet18', 'resnet34', 'resnet50', 'resnet101','resnet152'")
 
     global args, device
     args = parser.parse_args()