Skip to content
Permalink
Browse files

[Critical] use nn.ModuleList() to create layers or else it could only…

… do foward works!
  • Loading branch information...
RobinDong committed Sep 4, 2019
1 parent 1a87991 commit 9d14ed9eff145d4afaade7947b3189c4b534a97d
Showing with 9 additions and 5 deletions.
  1. +9 −5 ssd_mobilenet.py
@@ -39,7 +39,7 @@ def __init__(self, phase, size, extras, head, num_classes):
self.size = size

# SSD network
self.backbone = mobilenet.MobileNetV2(num_classes=num_classes, width_mult=1.0)
self.backbone = nn.ModuleList(mobilenet.MobileNetV2(num_classes=num_classes, width_mult=1.0).features)
self.norm = L2Norm(96, 20)
self.extras = nn.ModuleList(extras)

@@ -73,10 +73,14 @@ def forward(self, x):
loc = list()
conf = list()

_, endpoints = self.backbone(x)
sources.append(self.norm(endpoints[0]))
sources.append(endpoints[1])
x = endpoints[1]
for k in range(len(self.backbone)):
x = self.backbone[k](x)
if k in [13, 17]:
if len(sources) == 0:
s = self.norm(x)
sources.append(s)
else:
sources.append(x)

# apply extra layers and cache source layer outputs
for k, v in enumerate(self.extras):

0 comments on commit 9d14ed9

Please sign in to comment.
You can’t perform that action at this time.