未验证 提交 608c006b 编写于 作者: M Mansoor Nasir Cheema 提交者: GitHub

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
上级 624608fd
......@@ -223,12 +223,22 @@ class resnet(_fasterRCNN):
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)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册