| 26 | """ |
| 27 | |
| 28 | def __init__(self, phase, size, base, extras, head, num_classes): |
| 29 | super(SSD, self).__init__() |
| 30 | self.phase = phase |
| 31 | self.num_classes = num_classes |
| 32 | self.cfg = (coco, voc)[num_classes == 21] |
| 33 | self.priorbox = PriorBox(self.cfg) |
| 34 | self.priors = Variable(self.priorbox.forward(), volatile=True) |
| 35 | self.size = size |
| 36 | |
| 37 | # SSD network |
| 38 | self.vgg = nn.ModuleList(base) |
| 39 | # Layer learns to scale the l2 normalized features from conv4_3 |
| 40 | self.L2Norm = L2Norm(512, 20) |
| 41 | self.extras = nn.ModuleList(extras) |
| 42 | |
| 43 | self.loc = nn.ModuleList(head[0]) |
| 44 | self.conf = nn.ModuleList(head[1]) |
| 45 | |
| 46 | if phase == 'test': |
| 47 | self.softmax = nn.Softmax(dim=-1) |
| 48 | self.detect = Detect(num_classes, 0, 200, 0.01, 0.45) |
| 49 | |
| 50 | def forward(self, x): |
| 51 | """Applies network layers and ops on input image(s) x. |