| 213 | return x |
| 214 | |
| 215 | def init_params(self): |
| 216 | for m in self.modules(): |
| 217 | if isinstance(m, nn.Conv2d): |
| 218 | # init.kaiming_normal_(m.weight, mode='fan_out') |
| 219 | init.normal_(m.weight, std=0.01) |
| 220 | # init.xavier_normal_(m.weight) |
| 221 | if m.bias is not None: |
| 222 | init.constant_(m.bias, 0) |
| 223 | elif isinstance(m, nn.ConvTranspose2d): |
| 224 | # init.kaiming_normal_(m.weight, mode='fan_out') |
| 225 | init.normal_(m.weight, std=0.01) |
| 226 | # init.xavier_normal_(m.weight) |
| 227 | if m.bias is not None: |
| 228 | init.constant_(m.bias, 0) |
| 229 | elif isinstance(m, nn.BatchNorm2d): # NN.Batchnorm2d |
| 230 | init.constant_(m.weight, 1) |
| 231 | init.constant_(m.bias, 0) |
| 232 | elif isinstance(m, nn.Linear): |
| 233 | init.normal_(m.weight, std=0.01) |
| 234 | if m.bias is not None: |
| 235 | init.constant_(m.bias, 0) |
| 236 | |
| 237 | |
| 238 | class AO(nn.Module): |