Merge pull request #674 from mansoorcheema/master

Fix to use different resnet architectures based on input

Author: jwyang

lib/model/faster_rcnn/resnet.py

@@ -223,12 +223,22 @@ def __init__(self, classes, num_layers=101, pretrained=False, class_agnostic=Fal
     self.dout_base_model = 1024
     self.pretrained = pretrained
     self.class_agnostic = class_agnostic
+    self.num_layers = num_layers
 
     _fasterRCNN.__init__(self, classes, class_agnostic)
 
   def _init_modules(self):
     resnet = resnet101()
 
+    if self.num_layers == 18:
+        resnet = resnet18()
+    if self.num_layers == 34:
+        resnet = resnet34()     
+    if self.num_layers == 50:
+        resnet = resnet50()
+    if self.num_layers == 152:
+        resnet = resnet152()
+
     if self.pretrained == True:
       print("Loading pretrained weights from %s" %(self.model_path))
       state_dict = torch.load(self.model_path)