Lite-HRNet: A Lightweight High-Resolution Network

In plain words

Pose estimation and segmentation work best when a network keeps detailed, full-resolution features all the way through. The usual design for that, HRNet, is large. Small networks such as MobileNet or ShuffleNet are cheap, but it was unclear whether keeping high resolution helps them at all. Much of their remaining compute goes to mixing information across channels.

This paper starts by putting ShuffleNet-style blocks into a small HRNet. The main fix is a cheaper way to let channels talk to each other. Instead of a full mixing step, the network computes a set of weights from the features themselves and uses them to rescale the channels. The weights are gathered from every resolution the network is already processing, so they carry information across channels and resolutions at low cost.

The result is a network that is much cheaper than the large pose models while being more accurate than the other small networks the authors compare against. It is also applied to semantic segmentation with the same building blocks.

TL;DR

• Lite-HRNet is a lightweight version of HRNet built from shuffle blocks §3.1.
• Its key unit, conditional channel weighting, replaces the costly 1x1 convolutions with weights computed from all channels and all resolutions §3.2.
• On COCO val at 256x192, the authors report Lite-HRNet-30 at 67.2 AP with 1.8M parameters and 0.31 GFLOPs, without ImageNet pretraining tab. 3.
• The paper and the official repository agree on the core blocks, but the paper leaves many training details unstated. The repository fills them with specific values (see Paper vs code).
• No segmentation code is in the repository, so the segmentation results can only be checked against the paper.

Problem

High-resolution networks such as HRNet keep representations at high spatial resolution through the whole network, which helps position-sensitive tasks such as pose estimation and segmentation §1. Efficient networks built for classification (MobileNet, ShuffleNet) reduce cost mainly through depthwise convolutions. Depthwise convolutions do not mix information across channels, so these networks keep using pointwise 1x1 convolutions for that mixing. In a shuffle block, the two 1x1 convolutions cost more than the depthwise convolution, and the authors identify them as the bottleneck §3.2 tab. 2. The paper asks whether high resolution helps small models, and how to make a small high-resolution network cheap §1.

Method

Naive Lite-HRNet §3.1. The stem of Small HRNet (the authors' reference design) is changed so that its second 3x3 convolution is replaced by a shuffle block. All normal residual blocks, which are two 3x3 convolutions each, are replaced by shuffle blocks. The convolutions in the multi-resolution fusion are replaced by separable convolutions §3.1. The result is called naive Lite-HRNet. The authors use it as a baseline and show it already beats MobileNet, ShuffleNet and Small HRNet §1 fig. 4.

Conditional channel weighting §3.2. Each branch's channels are split into two halves §3.1. The first half passes through unchanged. The second half is rescaled element-wise by a weight map eq. 2:

\mathsf{Y}_s = \mathsf{W}_s \odot \mathsf{X}_s eq. 2

The weight map is computed in two parts.

1. Cross-resolution weights. Let there be s branches, with X_1 the highest resolution and X_s the lowest §3.2. Every branch except the lowest is average-pooled to the lowest branch's size. The pooled maps are concatenated with the lowest branch. A 1x1 convolution, a ReLU, another 1x1 convolution and a sigmoid produce s weight partitions, one per branch; the partitions for the higher branches are upsampled back to their sizes eq. 3 eq. 4:

   (\mathsf{W}_1, \dots, \mathsf{W}_s) = \mathcal{H}_s(\mathsf{X}_1, \dots, \mathsf{X}_s) eq. 3

   Each weight is then computed from the pooled channels of all resolutions at the same position, which lets the weight act as a bridge between channels and resolutions §3.2. The authors state that this scheme plays the same role as a 1x1 convolution for exchanging information across channels §3.2.

2. Spatial weights. For each branch, a single weight vector is computed from all positions of that branch. Global average pooling, a fully connected layer, ReLU, another fully connected layer and a sigmoid produce it §3.2 eq. 6. The vector is the same at every position.

The two weight types are combined in the conditional channel weighting block. At each position i, the output channels are the input channels multiplied by the cross-resolution weight at that position eq. 5:

\mathbf{y}_{si} = \mathbf{w}_{si} \odot \mathbf{x}_{si} eq. 5

The spatial weight is the same vector at every position eq. 6. The paper states that the two orders of applying the two weights perform similarly §4.3.

# Conditional channel weighting block  [§3.2] [eq. 2] [eq. 5]
# Input: list of s feature maps, one per resolution, (B, C_s, H_s, W_s)
# Output: list of s feature maps with the same shapes
def ccw_block(branches):
    # Split each branch into two channel halves  [§3.1]
    keep_halves = [first_half(b) for b in branches]
    work_halves = [second_half(b) for b in branches]
    # Cross-resolution weighting of the working halves  [§3.2] [eq. 3]
    work_halves = cross_resolution_weighting(work_halves)
    # Depthwise 3x3 convolution per branch; the paper does not state it here
    # [code: models/backbones/litehrnet.py#L120-L131]
    work_halves = [depthwise_3x3(h) for h in work_halves]
    # Spatial weighting per branch  [§3.2] [eq. 6]
    work_halves = [spatial_weighting(h) for h in work_halves]
    # Concatenate kept and working halves, then shuffle channels  [§3.1]
    outputs = [concat(k, w) for k, w in zip(keep_halves, work_halves)]
    return [channel_shuffle(o, groups=2) for o in outputs]

# Cross-resolution weight computation  [§3.2] [eq. 3] [eq. 4]
# Input: working halves X_1..X_s, X_1 is the highest resolution
# Output: each branch rescaled by its own weight map
def cross_resolution_weighting(halves):
    # Pool all branches except the lowest to the lowest one's size  [§3.2]
    lowest_size = size_of(halves[-1])
    pooled = [adaptive_avg_pool(h, lowest_size) for h in halves[:-1]]
    # Concatenate pooled branches with the lowest branch  [§3.2]
    all_channels = concat_channels(pooled + [halves[-1]])
    # 1x1 conv, ReLU, 1x1 conv, sigmoid  [eq. 4]
    weights = sigmoid(conv1x1(relu(conv1x1(all_channels))))
    # Split weights back per branch  [§3.2]
    weight_maps = split_channels(weights)
    # Upsample each weight map to its branch size and multiply  [§3.2]
    return [
        h * upsample_nearest(w, size_of(h))
        for h, w in zip(halves, weight_maps)
    ]

Key equations, as written in the paper:

(\mathsf{X}'_1, \dots, \mathsf{X}_s) \rightarrow \text{Conv} \rightarrow \text{ReLU} \rightarrow \text{Conv} \rightarrow \text{sigmoid} \rightarrow (\mathsf{W}'_1, \dots, \mathsf{W}'_s) eq. 4

\mathbf{w}_s = \mathcal{F}_s(\mathsf{X}_s) eq. 6

Spatial weighting detail litehrnet.py. The code uses 1x1 convolutions where the paper says fully connected layers; they are equivalent on pooled features. The reduction ratio of 4 is set in code litehrnet.py.

Lite module §3.2 tab. 1. One Lite-HRNet module consists of two conditional channel weighting blocks followed by one multi-resolution fusion block tab. 1 caption. The main body has three stages. Stage 2 has two branches, stage 3 has three and stage 4 has four. Branch channels are C, 2C, 4C and 8C §3.2. The table lists 40, 80, 160 and 320 channels, so C is 40 tab. 1. The number of modules is 2, 4 and 2 in the stages of Lite-HRNet-18, and 3, 8 and 3 in Lite-HRNet-30 tab. 1.

Stem §3.2 tab. 1. The stem has one 3x3 convolution with stride 2 and a shuffle block. It produces a 64x64 output for a 256x256 input tab. 1. The code's stem block downsamples itself (see Paper vs code).

Multi-resolution fusion §3.1 litehrnet.py. Each output branch sums resampled versions of all input branches, followed by ReLU. Upsampling is nearest-neighbour after a 1x1 convolution. Downsampling is a stack of stride-2 depthwise 3x3 and 1x1 convolutions. The paper does not describe this in detail.

Pose head and loss §4.1 litehrnet_18_coco_256x192.py. The network estimates K heatmaps. The head is a simple head with a 1x1 final convolution and no deconvolution. Training uses an MSE loss over heatmaps with per-joint target weights.

Hyperparameters

Paper values are cited to the paper; code values are cited to the official repository at code_commit. "not stated" means the paper gives no value.

Results

Author-reported. Setup: COCO keypoints (train2017, val2017, test-dev2017), MPII, Cityscapes §4.1 §4.4. Lite-HRNet entries are trained without ImageNet pretraining ("N" in the pretrain column). The pretrained baselines are reported as the authors give them. Parameter and FLOP counts are for the pose network, not the person detector tab. 3.

COCO val, tab. 3

COCO test-dev, tab. 4 (input 384x288 for the networks listed; large networks in the table are not repeated here)

The large networks in tab. 4 range from 63.1 AP (Mask-RCNN) to 76.2 AP (DARK). The paper says Lite-HRNet-30 is "significantly better" than the small networks and outperforms Mask-RCNN, G-RMI and Integral Pose Regression, while trailing some large networks §4.2.

MPII val, PCKh@0.5, tab. 5 (input 256x256; GFLOPs at that input)

Ablation: 1x1 convolution vs conditional channel weighting, tab. 6

Ablation: spatial vs cross-resolution weights, tab. 7

Semantic segmentation, Cityscapes, tab. 8 (trained from scratch, 512x1024 crops, no test-time tricks §4.4)

The paper compares these with hand-crafted and NAS-based methods in the same table, and says SwiftNet-RN18 is comparable at 104 GFLOPs §4.4 tab. 8. Lite-HRNet-18's 72.8 test mIoU is the value quoted in the text §4.4.

Paper vs code

Code: HRNet/Lite-HRNet at commit 7b9049d264fa40402a27d1f175deff3b46a6b91b (HEAD of the default branch hrnet, cloned at depth 1). The repository contains pose-estimation configs and the Lite-HRNet backbone. It has no segmentation code. The GitHub API shows its last push on 2022-06-03, so it is a later snapshot than the paper.

• Batch size. Paper: 32 per GPU on 8 V100 GPUs §4.1. Code: 64 per GPU for the 256x192 COCO configs litehrnet_18_coco_256x192.py litehrnet_30_coco_256x192.py; 32 per GPU for the 384x288 COCO configs litehrnet_18_coco_384x288.py and the MPII configs litehrnet_18_mpii_256x256.py. The configs do not state the total batch on 8 GPUs, but the per-GPU values differ from the paper's.
• Schedule. Paper: no epochs, learning-rate schedule or warmup for pose §4.1. Code: 210 epochs, step decay at 170 and 200, 500-iteration linear warmup litehrnet_18_coco_256x192.py.
• Heatmap sigma. Paper: not stated. Code: sigma 2 at 256x192 and 3 at 384x288 litehrnet_18_coco_256x192.py litehrnet_18_coco_384x288.py.
• Weight decay. Paper: stated only for segmentation (5e-4) §4.4. Code: the pose optimizer sets none litehrnet_18_coco_256x192.py.
• Channel-reduction ratios. Paper: not stated §3.2. Code: 8 for the cross-resolution weights litehrnet_18_coco_256x192.py and 4 for the spatial weights litehrnet.py.
• Normalization and activations in the weight paths. Paper: eq. 4 lists convolution, ReLU, convolution, sigmoid eq. 4. Code: the cross-resolution weight convolutions and the depthwise convolutions use batch normalization litehrnet.py litehrnet.py. The depthwise convolution has no activation litehrnet.py.
• Order of operations. Paper: not stated; it says a different arrangement performs similarly §4.3. Code: cross-resolution weighting, then depthwise convolution, then spatial weighting, applied to the second channel half only litehrnet.py.
• Upsampling. Paper: "upsampled", mode not stated. Code: nearest litehrnet.py.
• Stem. Paper: one 3x3 stride-2 convolution and a shuffle block tab. 1. Code: the stem block itself downsamples by 2, through a stride-2 depthwise convolution in the second half and a stride-2 branch on the first half litehrnet.py litehrnet.py.
• Augmentation. The paper's rotation range and scale range match the code §4.1 litehrnet_18_coco_256x192.py. Flip and half-body probabilities are in the code only.
• Test-time decoding. Paper: post-Gaussian filter and quarter offset §4.1. Code: flip test, modulate kernel 11, shift-heatmap flag set litehrnet_18_coco_256x192.py. We did not read the decoding function that consumes these flags.
• Segmentation. The repository has no segmentation code, so the segmentation results in tab. 8 cannot be checked against code.

Pitfalls

• Pretraining is not equal across rows. Lite-HRNet is trained without ImageNet pretraining. Several large and small baselines in tab. 3 are pretrained (Y). Compare the pretrain column before reading the AP gaps as architectural gains tab. 3.
• FLOP counts depend on input size. tab. 1 and tab. 5 use 256x256 inputs. tab. 3 uses 256x192 for COCO. Compare FLOPs at the same input size.
• Pose training details are missing from the paper. Epoch count, learning rate schedule, warmup and heatmap sigma are not stated in the paper. The code's values are the only concrete ones and should be marked as code-derived.
• Batch size differs. The paper says 32 per GPU. The 256x192 configs use 64 per GPU, so effective batch size on 8 GPUs will be twice what the paper describes if the configs are used as they are.
• The channel-reduction ratios are hidden hyperparameters. The paper gives the operations but not the ratios §3.2. The code uses 8 and 4. Changing them changes parameter count and FLOPs.
• Batch normalization is not in the paper's equations. Add it to the weight paths if you follow the code. Leaving it out will not match the code's training behaviour.
• Nearest-neighbour upsampling is a choice. The paper says only "upsampled". The code uses nearest.
• Stem downsampling is not explicit in the paper's text. Table 1 implies a 64x64 output for a 256x256 input, which the code's stem produces. Build the stem from the code, not from the text.
• Segmentation numbers have no code to check. Reproducing tab. 8 needs a segmentation pipeline the repository does not contain.
• Complexity savings are per example. The paper's 80% reduction (for one shape pair, 64x64x40 and 32x32x80) comes from its own example §1. The per-operation FLOP counts are for that example shape tab. 2.
• Accuracy depends on the training recipe. The paper does not state the schedule or epochs for pose. Reproducing the 67.2 AP needs the code's values, but the paper does not confirm that they produced it.

Lineage

• Builds on: HRNet, the high-resolution network with parallel multi-resolution streams §3.1 (Wang et al., TPAMI 2019, cited as Wang-TPAMI-HRNet-2019 in the paper). The paper uses the Small HRNet variant as its reference design §3.1.
• Builds on: ShuffleNet V2 shuffle blocks §3.1 (Ma et al., ECCV 2018).
• Related weight-generation methods: conditional convolutions, dynamic convolutions, dynamic filter networks and squeeze-and-excitation §2 [§3.2 Connection paragraph]. The paper's weights are used to exchange information across channels and resolutions, not to increase model capacity §3.2.
• Compared against: MobileNetV2, MobileNetV3, ShuffleNetV2, Small HRNet-W16, DY-MobileNetV2 and DY-ReLU tab. 3 tab. 5.
• Superseded by: not stated in content already in this wiki.