| 261 | |
| 262 | |
| 263 | class HighResolutionNet(nn.Module): |
| 264 | |
| 265 | def __init__(self, **kwargs): |
| 266 | extra = cfg.MODEL.OCR_EXTRA |
| 267 | super(HighResolutionNet, self).__init__() |
| 268 | |
| 269 | # stem net |
| 270 | self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, |
| 271 | bias=False) |
| 272 | self.bn1 = Norm2d(64, momentum=BN_MOMENTUM) |
| 273 | self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, |
| 274 | bias=False) |
| 275 | self.bn2 = Norm2d(64, momentum=BN_MOMENTUM) |
| 276 | self.relu = nn.ReLU(inplace=relu_inplace) |
| 277 | |
| 278 | self.stage1_cfg = extra['STAGE1'] |
| 279 | num_channels = self.stage1_cfg['NUM_CHANNELS'][0] |
| 280 | block = blocks_dict[self.stage1_cfg['BLOCK']] |
| 281 | num_blocks = self.stage1_cfg['NUM_BLOCKS'][0] |
| 282 | self.layer1 = self._make_layer(block, 64, num_channels, num_blocks) |
| 283 | stage1_out_channel = block.expansion*num_channels |
| 284 | |
| 285 | self.stage2_cfg = extra['STAGE2'] |
| 286 | num_channels = self.stage2_cfg['NUM_CHANNELS'] |
| 287 | block = blocks_dict[self.stage2_cfg['BLOCK']] |
| 288 | num_channels = [num_channels[i] * block.expansion |
| 289 | for i in range(len(num_channels))] |
| 290 | self.transition1 = self._make_transition_layer( |
| 291 | [stage1_out_channel], num_channels) |
| 292 | self.stage2, pre_stage_channels = self._make_stage( |
| 293 | self.stage2_cfg, num_channels) |
| 294 | |
| 295 | self.stage3_cfg = extra['STAGE3'] |
| 296 | num_channels = self.stage3_cfg['NUM_CHANNELS'] |
| 297 | block = blocks_dict[self.stage3_cfg['BLOCK']] |
| 298 | num_channels = [num_channels[i] * block.expansion |
| 299 | for i in range(len(num_channels))] |
| 300 | self.transition2 = self._make_transition_layer( |
| 301 | pre_stage_channels, num_channels) |
| 302 | self.stage3, pre_stage_channels = self._make_stage( |
| 303 | self.stage3_cfg, num_channels) |
| 304 | |
| 305 | self.stage4_cfg = extra['STAGE4'] |
| 306 | num_channels = self.stage4_cfg['NUM_CHANNELS'] |
| 307 | block = blocks_dict[self.stage4_cfg['BLOCK']] |
| 308 | num_channels = [num_channels[i] * block.expansion |
| 309 | for i in range(len(num_channels))] |
| 310 | self.transition3 = self._make_transition_layer( |
| 311 | pre_stage_channels, num_channels) |
| 312 | self.stage4, pre_stage_channels = self._make_stage( |
| 313 | self.stage4_cfg, num_channels, multi_scale_output=True) |
| 314 | |
| 315 | self.high_level_ch = np.int(np.sum(pre_stage_channels)) |
| 316 | |
| 317 | def _make_transition_layer( |
| 318 | self, num_channels_pre_layer, num_channels_cur_layer): |
| 319 | num_branches_cur = len(num_channels_cur_layer) |
| 320 | num_branches_pre = len(num_channels_pre_layer) |