RegNet model. Paper: https://arxiv.org/abs/2003.13678 Original Impl: https://github.com/facebookresearch/pycls/blob/master/pycls/models/regnet.py and refer to: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/regnet.py
| 217 | return x |
| 218 | |
| 219 | class RegNet(nn.Module): |
| 220 | """RegNet model. |
| 221 | Paper: https://arxiv.org/abs/2003.13678 |
| 222 | Original Impl: https://github.com/facebookresearch/pycls/blob/master/pycls/models/regnet.py |
| 223 | and refer to: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/regnet.py |
| 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(): |
| 266 | x = layer(x) |
| 267 | return x |
| 268 | |
| 269 | @staticmethod |
| 270 | def _build_stage_info(cfg: dict): |
| 271 | wa, w0, wm, d = cfg["wa"], cfg["w0"], cfg["wm"], cfg["depth"] |
| 272 | widths, num_stages = generate_width_depth(wa, w0, wm, d) |
| 273 | |
| 274 | stage_widths, stage_depths = np.unique(widths, return_counts=True) |
| 275 | stage_groups = [cfg['group_w'] for _ in range(num_stages)] |
| 276 | stage_widths, stage_groups = adjust_width_groups_comp(stage_widths, stage_groups) |