def __init__(self, extra, norm_eval=True, zero_init_residual=False, frozen_stages=(- 1)):
    super(HighResolutionNet, self).__init__()
    self.norm_eval = norm_eval
    self.frozen_stages = frozen_stages
    self.zero_init_residual = zero_init_residual
    self.extra = extra
    self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
    self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
    self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
    self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
    self.relu = nn.ReLU(inplace=True)
    self.stage1_cfg = self.extra['stage1']
    num_channels = self.stage1_cfg['num_channels'][0]
    block_type = self.stage1_cfg['block']
    num_blocks = self.stage1_cfg['num_blocks'][0]
    block = blocks_dict[block_type]
    stage1_out_channels = (num_channels * block.expansion)
    self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
    self.stage2_cfg = self.extra['stage2']
    num_channels = self.stage2_cfg['num_channels']
    block_type = self.stage2_cfg['block']
    block = blocks_dict[block_type]
    num_channels = [(num_channels[i] * block.expansion) for i in range(len(num_channels))]
    self.transition1 = self._make_transition_layer([stage1_out_channels], num_channels)
    (self.stage2, pre_stage_channels) = self._make_stage(self.stage2_cfg, num_channels)
    self.stage3_cfg = self.extra['stage3']
    num_channels = self.stage3_cfg['num_channels']
    block_type = self.stage3_cfg['block']
    block = blocks_dict[block_type]
    num_channels = [(num_channels[i] * block.expansion) for i in range(len(num_channels))]
    self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
    (self.stage3, pre_stage_channels) = self._make_stage(self.stage3_cfg, num_channels)
    self.stage4_cfg = self.extra['stage4']
    num_channels = self.stage4_cfg['num_channels']
    block_type = self.stage4_cfg['block']
    block = blocks_dict[block_type]
    num_channels = [(num_channels[i] * block.expansion) for i in range(len(num_channels))]
    self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
    (self.stage4, pre_stage_channels) = self._make_stage(self.stage4_cfg, num_channels)
    self._frozen_stages()
def __init__(self, extra, norm_eval=True, zero_init_residual=False, frozen_stages=(- 1)):
    super(HighResolutionNet, self).__init__()
    self.norm_eval = norm_eval
    self.frozen_stages = frozen_stages
    self.zero_init_residual = zero_init_residual
    self.extra = extra
    self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
    self.bn1 = BatchNorm2d(64, momentum=BN_MOMENTUM)
    self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
    self.bn2 = BatchNorm2d(64, momentum=BN_MOMENTUM)
    self.relu = nn.ReLU(inplace=True)
    self.stage1_cfg = self.extra['stage1']
    num_channels = self.stage1_cfg['num_channels'][0]
    block_type = self.stage1_cfg['block']
    num_blocks = self.stage1_cfg['num_blocks'][0]
    block = blocks_dict[block_type]
    stage1_out_channels = (num_channels * block.expansion)
    self.layer1 = self._make_layer(block, 64, num_channels, num_blocks)
    self.stage2_cfg = self.extra['stage2']
    num_channels = self.stage2_cfg['num_channels']
    block_type = self.stage2_cfg['block']
    block = blocks_dict[block_type]
    num_channels = [(num_channels[i] * block.expansion) for i in range(len(num_channels))]
    self.transition1 = self._make_transition_layer([stage1_out_channels], num_channels)
    (self.stage2, pre_stage_channels) = self._make_stage(self.stage2_cfg, num_channels)
    self.stage3_cfg = self.extra['stage3']
    num_channels = self.stage3_cfg['num_channels']
    block_type = self.stage3_cfg['block']
    block = blocks_dict[block_type]
    num_channels = [(num_channels[i] * block.expansion) for i in range(len(num_channels))]
    self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
    (self.stage3, pre_stage_channels) = self._make_stage(self.stage3_cfg, num_channels)
    self.stage4_cfg = self.extra['stage4']
    num_channels = self.stage4_cfg['num_channels']
    block_type = self.stage4_cfg['block']
    block = blocks_dict[block_type]
    num_channels = [(num_channels[i] * block.expansion) for i in range(len(num_channels))]
    self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
    (self.stage4, pre_stage_channels) = self._make_stage(self.stage4_cfg, num_channels)
    self._frozen_stages()