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 | def __init__(self, |
| 21 | num_branches, |
| 22 | blocks, |
| 23 | num_blocks, |
| 24 | in_channels, |
| 25 | num_channels, |
| 26 | multiscale_output=True, |
| 27 | with_cp=False, |
| 28 | conv_cfg=None, |
| 29 | norm_cfg=dict(type='BN'), |
| 30 | block_init_cfg=None, |
| 31 | init_cfg=None): |
| 32 | super(HRModule, self).__init__(init_cfg) |
| 33 | self.block_init_cfg = block_init_cfg |
| 34 | self._check_branches(num_branches, num_blocks, in_channels, |
| 35 | num_channels) |
| 36 | |
| 37 | self.in_channels = in_channels |
| 38 | self.num_branches = num_branches |
| 39 | |
| 40 | self.multiscale_output = multiscale_output |
| 41 | self.norm_cfg = norm_cfg |
| 42 | self.conv_cfg = conv_cfg |
| 43 | self.with_cp = with_cp |
| 44 | self.branches = self._make_branches(num_branches, blocks, num_blocks, |
| 45 | num_channels) |
| 46 | self.fuse_layers = self._make_fuse_layers() |
| 47 | self.relu = nn.ReLU(inplace=False) |
| 48 | |
| 49 | def _check_branches(self, num_branches, num_blocks, in_channels, |
| 50 | num_channels): |
| 51 | if num_branches != len(num_blocks): |
| 52 | error_msg = f'NUM_BRANCHES({num_branches}) ' \ |
| 53 | f'!= 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}) ' \ |
| 58 | f'!= NUM_CHANNELS({len(num_channels)})' |
| 59 | raise ValueError(error_msg) |
| 60 | |
| 61 | if num_branches != len(in_channels): |
| 62 | error_msg = f'NUM_BRANCHES({num_branches}) ' \ |
| 63 | f'!= NUM_INCHANNELS({len(in_channels)})' |
| 64 | raise ValueError(error_msg) |
| 65 | |
| 66 | def _make_one_branch(self, |
| 67 | branch_index, |
| 68 | block, |
| 69 | num_blocks, |
| 70 | num_channels, |
| 71 | stride=1): |