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__( |
| 22 | self, |
| 23 | num_branches, |
| 24 | blocks, |
| 25 | num_blocks, |
| 26 | in_channels, |
| 27 | num_channels, |
| 28 | multiscale_output=True, |
| 29 | with_cp=False, |
| 30 | conv_cfg=None, |
| 31 | norm_cfg=dict(type="BN", requires_grad=True), |
| 32 | block_init_cfg=None, |
| 33 | init_cfg=None, |
| 34 | ): |
| 35 | super(HRModule, self).__init__(init_cfg) |
| 36 | self.block_init_cfg = block_init_cfg |
| 37 | self._check_branches(num_branches, num_blocks, in_channels, num_channels) |
| 38 | |
| 39 | self.in_channels = in_channels |
| 40 | self.num_branches = num_branches |
| 41 | |
| 42 | self.multiscale_output = multiscale_output |
| 43 | self.norm_cfg = norm_cfg |
| 44 | self.conv_cfg = conv_cfg |
| 45 | self.with_cp = with_cp |
| 46 | self.branches = self._make_branches(num_branches, blocks, num_blocks, 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, num_channels): |
| 51 | """Check branches configuration.""" |
| 52 | if num_branches != len(num_blocks): |
| 53 | error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_BLOCKS({len(num_blocks)})" |
| 54 | raise ValueError(error_msg) |
| 55 | |
| 56 | if num_branches != len(num_channels): |
| 57 | error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_CHANNELS({len(num_channels)})" |
| 58 | raise ValueError(error_msg) |
| 59 | |
| 60 | if num_branches != len(in_channels): |
| 61 | error_msg = f"NUM_BRANCHES({num_branches}) <> NUM_INCHANNELS({len(in_channels)})" |
| 62 | raise ValueError(error_msg) |
| 63 | |
| 64 | def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1): |
| 65 | """Build one branch.""" |
| 66 | downsample = None |
| 67 | if stride != 1 or self.in_channels[branch_index] != num_channels[branch_index] * block.expansion: |
| 68 | downsample = nn.Sequential( |
| 69 | build_conv_layer( |
| 70 | self.conv_cfg, |
| 71 | self.in_channels[branch_index], |