Add option in eval.py to use different resnets. (Must match with what you use to preprocess.)

Author: ruotianluo

dataloaderraw.py

@@ -22,12 +22,6 @@
 from misc.resnet_utils import myResnet
 import misc.resnet as resnet
 
-resnet = resnet.resnet101()
-resnet.load_state_dict(torch.load('./data/imagenet_weights/resnet101.pth'))
-my_resnet = myResnet(resnet)
-my_resnet.cuda()
-my_resnet.eval()
-
 class DataLoaderRaw():
 
     def __init__(self, opt):
@@ -38,6 +32,16 @@ def __init__(self, opt):
         self.batch_size = opt.get('batch_size', 1)
         self.seq_per_img = 1
 
+        # Load resnet
+        self.cnn_model = opt.get('cnn_model', 'resnet101')
+        resnet = getattr(resnet, self.cnn_model)()
+        resnet.load_state_dict(torch.load('./data/imagenet_weights/'+self.cnn_model+'.pth'))
+        self.my_resnet = myResnet(resnet)
+        self.my_resnet.cuda()
+        self.my_resnet.eval()
+
+
+
         # load the json file which contains additional information about the dataset
         print('DataLoaderRaw loading images from folder: ', self.folder_path)
 
@@ -106,7 +110,7 @@ def get_batch(self, split, batch_size=None):
             img = img.astype('float32')/255.0
             img = torch.from_numpy(img.transpose([2,0,1])).cuda()
             img = Variable(preprocess(img), volatile=True)
-            tmp_fc, tmp_att = my_resnet(img)
+            tmp_fc, tmp_att = self.my_resnet(img)
 
             fc_batch[i] = tmp_fc.data.cpu().float().numpy()
             att_batch[i] = tmp_att.data.cpu().float().numpy()

eval.py

@@ -23,6 +23,8 @@
 # Input paths
 parser.add_argument('--model', type=str, default='',
                 help='path to model to evaluate')
+parser.add_argument('--cnn_model', type=str,  default='resnet101',
+                help='resnet101, resnet152')
 parser.add_argument('--infos_path', type=str, default='',
                 help='path to infos to evaluate')
 # Basic options
@@ -108,7 +110,8 @@
 else:
   loader = DataLoaderRaw({'folder_path': opt.image_folder, 
                             'coco_json': opt.coco_json,
-                            'batch_size': opt.batch_size})
+                            'batch_size': opt.batch_size,
+                            'cnn_model': opt.cnn_model})
   loader.ix_to_word = infos['vocab']