(self,
cfg: dict,
in_c: int = 3,
num_classes: int = 1000,
zero_init_last_bn: bool = True)
| 224 | """ |
| 225 | |
| 226 | def __init__(self, |
| 227 | cfg: dict, |
| 228 | in_c: int = 3, |
| 229 | num_classes: int = 1000, |
| 230 | zero_init_last_bn: bool = True): |
| 231 | super(RegNet, self).__init__() |
| 232 | |
| 233 | # RegStem |
| 234 | stem_c = cfg["stem_width"] |
| 235 | self.stem = ConvBNAct(in_c, out_c=stem_c, kernel_s=3, stride=2, padding=1) |
| 236 | |
| 237 | # build stages |
| 238 | input_channels = stem_c |
| 239 | stage_info = self._build_stage_info(cfg) |
| 240 | for i, stage_args in enumerate(stage_info): |
| 241 | stage_name = "s{}".format(i + 1) |
| 242 | self.add_module(stage_name, RegStage(in_c=input_channels, **stage_args)) |
| 243 | input_channels = stage_args["out_c"] |
| 244 | |
| 245 | # RegHead |
| 246 | self.head = RegHead(in_unit=input_channels, out_unit=num_classes) |
| 247 | |
| 248 | # initial weights |
| 249 | for m in self.modules(): |
| 250 | if isinstance(m, nn.Conv2d): |
| 251 | nn.init.kaiming_uniform_(m.weight, mode="fan_out", nonlinearity='relu') |
| 252 | elif isinstance(m, nn.BatchNorm2d): |
| 253 | nn.init.ones_(m.weight) |
| 254 | nn.init.zeros_(m.bias) |
| 255 | elif isinstance(m, nn.Linear): |
| 256 | nn.init.normal_(m.weight, mean=0.0, std=0.01) |
| 257 | nn.init.zeros_(m.bias) |
| 258 | |
| 259 | if zero_init_last_bn: |
| 260 | for m in self.modules(): |
| 261 | if hasattr(m, "zero_init_last_bn"): |
| 262 | m.zero_init_last_bn() |
| 263 | |
| 264 | def forward(self, x: Tensor) -> Tensor: |
| 265 | for layer in self.children(): |
no test coverage detected