Skip to content

Commit 608c006

Browse files
Fix to use different resnet architectures based on input
Current code always uses resnet 101 architecture for all resnet based architectures while ignoring the arguments for current architecture
1 parent 624608f commit 608c006

File tree

1 file changed

+10
-0
lines changed

1 file changed

+10
-0
lines changed

lib/model/faster_rcnn/resnet.py

+10
Original file line numberDiff line numberDiff line change
@@ -223,12 +223,22 @@ def __init__(self, classes, num_layers=101, pretrained=False, class_agnostic=Fal
223223
self.dout_base_model = 1024
224224
self.pretrained = pretrained
225225
self.class_agnostic = class_agnostic
226+
self.num_layers = num_layers
226227

227228
_fasterRCNN.__init__(self, classes, class_agnostic)
228229

229230
def _init_modules(self):
230231
resnet = resnet101()
231232

233+
if self.num_layers == 18:
234+
resnet = resnet18()
235+
if self.num_layers == 34:
236+
resnet = resnet34()
237+
if self.num_layers == 50:
238+
resnet = resnet50()
239+
if self.num_layers == 152:
240+
resnet = resnet152()
241+
232242
if self.pretrained == True:
233243
print("Loading pretrained weights from %s" %(self.model_path))
234244
state_dict = torch.load(self.model_path)

0 commit comments

Comments
 (0)