| 14 | } |
| 15 | |
| 16 | class VGG(nn.Module): |
| 17 | def __init__(self, features, num_classes=1000, init_weights=False): |
| 18 | super(VGG, self).__init__() |
| 19 | self.features = features |
| 20 | self.classifier = nn.Sequential( |
| 21 | nn.Linear(512*7*7, 4096), |
| 22 | nn.ReLU(True), |
| 23 | nn.Dropout(p=0.5), |
| 24 | nn.Linear(4096, 4096), |
| 25 | nn.ReLU(True), |
| 26 | nn.Dropout(p=0.5), |
| 27 | nn.Linear(4096, num_classes) |
| 28 | ) |
| 29 | if init_weights: |
| 30 | self._initialize_weights() |
| 31 | |
| 32 | def forward(self, x): |
| 33 | # N x 3 x 224 x 224 |
| 34 | x = self.features(x) |
| 35 | # N x 512 x 7 x 7 |
| 36 | x = torch.flatten(x, start_dim=1) |
| 37 | # N x 512*7*7 |
| 38 | x = self.classifier(x) |
| 39 | return x |
| 40 | |
| 41 | def _initialize_weights(self): |
| 42 | for m in self.modules(): |
| 43 | if isinstance(m, nn.Conv2d): |
| 44 | # nn.init.kaiming_normal_(m.weights, mode='fan_out', nonlinearity='relu') |
| 45 | nn.init.xavier_uniform_(m.weight) |
| 46 | if m.bias is not None: |
| 47 | nn.init.constant_(m.bias, 0) |
| 48 | elif isinstance(m, nn.Linear): |
| 49 | nn.init.xavier_uniform_(m.weight) |
| 50 | # nn.init.normal_(m.weight, 0, 0.01) |
| 51 | nn.init.constant_(m.bias, 0) |
| 52 | |
| 53 | # 采用写配置文件的方式，将网络对应层的类型及其参数写入到sequential容器中 |
| 54 | def make_features(cfg: list): |