High-Resolution Module for HRNet. In this module, every branch has 4 BasicBlocks/Bottlenecks. Fusion/Exchange is in this module.
| 12 | |
| 13 | |
| 14 | class HRModule(BaseModule): |
| 15 | """High-Resolution Module for HRNet. |
| 16 | |
| 17 | In this module, every branch has 4 BasicBlocks/Bottlenecks. Fusion/Exchange |
| 18 | is in this module. |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, |
| 22 | num_branches, |
| 23 | blocks, |
| 24 | num_blocks, |
| 25 | in_channels, |
| 26 | num_channels, |
| 27 | multiscale_output=True, |
| 28 | with_cp=False, |
| 29 | conv_cfg=None, |
| 30 | norm_cfg=dict(type='BN', requires_grad=True), |
| 31 | block_init_cfg=None, |
| 32 | init_cfg=None): |
| 33 | super(HRModule, self).__init__(init_cfg) |
| 34 | self.block_init_cfg = block_init_cfg |
| 35 | self._check_branches(num_branches, num_blocks, in_channels, |
| 36 | num_channels) |
| 37 | |
| 38 | self.in_channels = in_channels |
| 39 | self.num_branches = num_branches |
| 40 | |
| 41 | self.multiscale_output = multiscale_output |
| 42 | self.norm_cfg = norm_cfg |
| 43 | self.conv_cfg = conv_cfg |
| 44 | self.with_cp = with_cp |
| 45 | self.branches = self._make_branches(num_branches, blocks, num_blocks, |
| 46 | num_channels) |
| 47 | self.fuse_layers = self._make_fuse_layers() |
| 48 | self.relu = nn.ReLU(inplace=False) |
| 49 | |
| 50 | def _check_branches(self, num_branches, num_blocks, in_channels, |
| 51 | num_channels): |
| 52 | """Check branches configuration.""" |
| 53 | if num_branches != len(num_blocks): |
| 54 | error_msg = f'NUM_BRANCHES({num_branches}) <> NUM_BLOCKS(' \ |
| 55 | f'{len(num_blocks)})' |
| 56 | raise ValueError(error_msg) |
| 57 | |
| 58 | if num_branches != len(num_channels): |
| 59 | error_msg = f'NUM_BRANCHES({num_branches}) <> NUM_CHANNELS(' \ |
| 60 | f'{len(num_channels)})' |
| 61 | raise ValueError(error_msg) |
| 62 | |
| 63 | if num_branches != len(in_channels): |
| 64 | error_msg = f'NUM_BRANCHES({num_branches}) <> NUM_INCHANNELS(' \ |
| 65 | f'{len(in_channels)})' |
| 66 | raise ValueError(error_msg) |
| 67 | |
| 68 | def _make_one_branch(self, |
| 69 | branch_index, |
| 70 | block, |
| 71 | num_blocks, |