(self, cfg, in_chans=3, num_classes=1000, global_pool='avg', drop_rate=0.0, head='classification')
| 507 | class HighResolutionNet(nn.Module): |
| 508 | |
| 509 | def __init__(self, cfg, in_chans=3, num_classes=1000, global_pool='avg', drop_rate=0.0, head='classification'): |
| 510 | super(HighResolutionNet, self).__init__() |
| 511 | self.num_classes = num_classes |
| 512 | self.drop_rate = drop_rate |
| 513 | |
| 514 | stem_width = cfg['STEM_WIDTH'] |
| 515 | self.conv1 = nn.Conv2d(in_chans, stem_width, kernel_size=3, stride=2, padding=1, bias=False) |
| 516 | self.bn1 = nn.BatchNorm2d(stem_width, momentum=_BN_MOMENTUM) |
| 517 | self.act1 = nn.ReLU(inplace=True) |
| 518 | self.conv2 = nn.Conv2d(stem_width, 64, kernel_size=3, stride=2, padding=1, bias=False) |
| 519 | self.bn2 = nn.BatchNorm2d(64, momentum=_BN_MOMENTUM) |
| 520 | self.act2 = nn.ReLU(inplace=True) |
| 521 | |
| 522 | self.stage1_cfg = cfg['STAGE1'] |
| 523 | num_channels = self.stage1_cfg['NUM_CHANNELS'][0] |
| 524 | block = blocks_dict[self.stage1_cfg['BLOCK']] |
| 525 | num_blocks = self.stage1_cfg['NUM_BLOCKS'][0] |
| 526 | self.layer1 = self._make_layer(block, 64, num_channels, num_blocks) |
| 527 | stage1_out_channel = block.expansion * num_channels |
| 528 | |
| 529 | self.stage2_cfg = cfg['STAGE2'] |
| 530 | num_channels = self.stage2_cfg['NUM_CHANNELS'] |
| 531 | block = blocks_dict[self.stage2_cfg['BLOCK']] |
| 532 | num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))] |
| 533 | self.transition1 = self._make_transition_layer([stage1_out_channel], num_channels) |
| 534 | self.stage2, pre_stage_channels = self._make_stage(self.stage2_cfg, num_channels) |
| 535 | |
| 536 | self.stage3_cfg = cfg['STAGE3'] |
| 537 | num_channels = self.stage3_cfg['NUM_CHANNELS'] |
| 538 | block = blocks_dict[self.stage3_cfg['BLOCK']] |
| 539 | num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))] |
| 540 | self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels) |
| 541 | self.stage3, pre_stage_channels = self._make_stage(self.stage3_cfg, num_channels) |
| 542 | |
| 543 | self.stage4_cfg = cfg['STAGE4'] |
| 544 | num_channels = self.stage4_cfg['NUM_CHANNELS'] |
| 545 | block = blocks_dict[self.stage4_cfg['BLOCK']] |
| 546 | num_channels = [num_channels[i] * block.expansion for i in range(len(num_channels))] |
| 547 | self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels) |
| 548 | self.stage4, pre_stage_channels = self._make_stage(self.stage4_cfg, num_channels, multi_scale_output=True) |
| 549 | |
| 550 | self.head = head |
| 551 | self.head_channels = None # set if _make_head called |
| 552 | if head == 'classification': |
| 553 | # Classification Head |
| 554 | self.num_features = 2048 |
| 555 | self.incre_modules, self.downsamp_modules, self.final_layer = self._make_head(pre_stage_channels) |
| 556 | self.global_pool, self.classifier = create_classifier( |
| 557 | self.num_features, self.num_classes, pool_type=global_pool) |
| 558 | elif head == 'incre': |
| 559 | self.num_features = 2048 |
| 560 | self.incre_modules, _, _ = self._make_head(pre_stage_channels, True) |
| 561 | else: |
| 562 | self.incre_modules = None |
| 563 | self.num_features = 256 |
| 564 | |
| 565 | curr_stride = 2 |
| 566 | # module names aren't actually valid here, hook or FeatureNet based extraction would not work |
no test coverage detected