(self, features, out_channels)
| 16 | class VGG(nn.Module): |
| 17 | ''' VGG model ''' |
| 18 | def __init__(self, features, out_channels): |
| 19 | super(VGG, self).__init__() |
| 20 | self.features = features |
| 21 | self.avgpool = nn.AdaptiveAvgPool2d((1,1)) |
| 22 | self.classifier = nn.Sequential( |
| 23 | nn.Dropout(), |
| 24 | nn.Linear(512, 512), |
| 25 | nn.ReLU(True), |
| 26 | nn.Dropout(), |
| 27 | nn.Linear(512, 512), |
| 28 | nn.ReLU(True), |
| 29 | nn.Linear(512, out_channels), |
| 30 | ) |
| 31 | |
| 32 | ''' Initialize weights ''' |
| 33 | for m in self.modules(): |
| 34 | if isinstance(m, nn.Conv2d): |
| 35 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
| 36 | m.weight.data.normal_(0, math.sqrt(2. / n)) |
| 37 | m.bias.data.zero_() |
| 38 | |
| 39 | |
| 40 | def forward(self, x): |
nothing calls this directly
no outgoing calls
no test coverage detected