Layout Matters: Optimizing AggPose Inference with OpenVINO

Background

AggPose is a model released in 2022 for infant pose estimation.

I remain somewhat skeptical about the motivation behind designing a dedicated architecture for this task¹. What exactly is the task-specific prior that makes infant pose estimation fundamentally different? Why not start from a strong general-purpose human pose estimation architecture and fine-tune it on infant data? And if this architecture were broadly competitive, it would be natural to ask why it was not evaluated more aggressively against SOTA methods for adult human pose estimation.

That said, given the actual project requirements, the practical question becomes how to optimize AggPose as much as possible. Our main goal is to run AggPose on hospital-side machines for offline GMA video processing, extract infant poses, and obtain de-identified infant pose data. Therefore, we focus on maximizing throughput rather than minimizing single-frame latency. We also assume that the hospital-side machines only have CPUs, so the optimization target is a single-CPU deployment scenario.

For now, we do not yet have the exact hardware specifications of the hospital-side machines. As a result, the current optimization work is mainly performed on the AMD Ryzen 7600X CPU available in our home lab. Once the target machine specifications are available, we can migrate and re-tune the optimization accordingly.

Model

TL;DR: SegFormer with an HRNet-style flavor.

The implementation is available in the PediaMedAI/AggPose² repository. The authors released two checkpoints:

• AggPose-L_256x192_COCO2017.pth: trained on the COCO 17 keypoint dataset, with 17 keypoints.
• AggPose-L_256x192_infantpose2022.pth: presumably trained on the authors’ InfantPose dataset, with 21 keypoints.

Because we do not have access to the corresponding infant dataset, it is difficult to perform calibration and distillation on infant data. Therefore, the current optimization work is based on the COCO-17 checkpoint.

AggPose follows the top-down pose estimation paradigm. In other words, it must be used together with a person detector. In this optimization work, however, we do not include the detector.

The overall architecture is similar to HRNet: different branches correspond to different feature-map resolutions.

However, AggPose replaces each channel map branch with a modified AViT module, short for Aggregation Vision Transformer. Each AViT block is structured as follows:

Compared with the Transformer Encoder in ViT, AggPose replaces MHSA (Multi-Head Self-Attention) with Eff-Attention³.

Optimization

Experiment Setup

The experiments were run on an AMD Ryzen 7600X with 6 cores and 12 threads.

Accuracy evaluation is mainly based on COCO 17. We use throughput as the primary performance metric, while latency is reported only as a secondary reference.

One important caveat is that the reported throughput does not represent end-to-end throughput in a real deployment. It does not include the overhead of video decoding, preprocessing, the person detector, or postprocessing.

The optimization steps below are presented in a slightly rearranged order and do not necessarily reflect the exact chronological order of experimentation.

Baseline

We first rewrote the original implementation into a more parameterized version. The baseline implementations use PyTorch, including eager mode and torch.compile, ONNX Runtime, and OpenVINO. The OpenVINO baseline is not manually tuned but uses Async Infer. All baselines use FP32 inference.

OpenVINO performs very well out of the box, which is why we chose it as the basis for further inference optimization.

Step 1: Quantization

We use NNCF, which is closely aligned with the OpenVINO ecosystem, for PTQ, or Post-Training Quantization. Since the CPU supports avx512_vnni, we quantize the model to int8.

The number of calibration samples needs to be reasonably large. I used 300 samples; using a much smaller calibration set caused severe accuracy degradation, with AP dropping to 0.

However, after quantization, throughput became worse rather than better. This issue is addressed later in Step 4.

Step 2: Static Batch

The model input is dynamic, but only along the batch_size dimension.

Specifically, the input shape is [B, 3, 256, 192], where B is the inference-time batch_size. For a given hardware platform, there should be an optimal batch_size that maximizes throughput.

First, we benchmarked the dynamic-shape model with different batch sizes:

Then we converted the high-throughput dynamic-shape candidates, namely batch sizes 4, 8, and 16, into static-shape models and benchmarked them again:

The best batch size is therefore 4.

Step 3: Pruning + Distillation

We noticed that AggPose stacks 40 blocks in AViT_3,1, which seems excessive. Earlier depths are at most 6 blocks. We therefore tried reducing the depth of this stage.

We tested five variants: D32, D24, D16, D12, and D8. Throughput was measured with Torch CPU eager mode.

We selected the D24 variant and distilled it from the original D40 model to recover part of the lost accuracy. The distillation objective follows Knowledge Distillation:

where:

• $H_{\text{T}}$ denotes the teacher model’s predicted heatmap.
• $H_{\text{S}}$ denotes the student model’s predicted heatmap.
• $T$ is the distillation temperature, usually set to 2.

The training data is the full COCO-17 training set. We mainly trained for two epochs, although one epoch may already be sufficient:

• In the first epoch, only stage 2 is fine-tuned. In other words, stage 0 and stage 1 are frozen, namely AViT 1,1, AViT 2,1, and AViT 1,2.
• In the second epoch, the full model is fine-tuned.

The following table shows full validation results on the COCO-17 validation set for the teacher model and the distilled student model:

As shown above, the distilled student model recovers most of the original model’s accuracy.

Step 4: Replacing Linear with 1x1 Conv2d

As mentioned earlier, after int8 quantization, throughput became even worse than the fp32 version. In this step, we locate and address the underlying issue.

Profiling

Using benchmark_app to profile the performance hot spots, we found that FakeQuantize and Reorder had unusually high runtime. The breakdown is as follows:

Locating the Problem

A rough inspection showed that the Reorder and FakeQuantize operators mainly appear in the mlp module, corresponding to the Mlp module in AggPose. We then used benchmark_app for a more detailed profile.

Based on the report exported by benchmark_app in average_counters mode, namely benchmark_average_counters_report.csv, and comparing it against the model structure file .xml, we can inspect the actually executed layers.

Take __module.block_passes.0.levels.0.blocks.1.mlp as an example. After removing operators that were not executed, that is, entries whose execStatus is Status.NOT_RUN, the actual layers in the mlp module are:

The per-layer runtime is shown below:

Reorder and FakeQuantize account for most of the runtime. It is also clear that the three Reorder nodes are introduced around GroupConvolution and FakeQuantize.

Source-Level Insight

Mapping the profile results back to the model source code, the issue is located in Mlp, which is the Mix-FFN module from SegFormer:

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

        self.apply(self._init_weights)

    def _init_weights(self, m):
        ...

    def forward(self, x, H, W):
        x = self.fc1(x)
        x = self.dwconv(x, H, W)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

A depth-wise convolution operation, dwconv, is inserted after fc1. The implementation is in DWConv:

class DWConv(nn.Module):
    def __init__(self, dim=768):
        super(DWConv, self).__init__()
        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

    def forward(self, x, H, W):
        B, N, C = x.shape
        x = x.transpose(1, 2).view(B, C, H, W)
        x = self.dwconv(x)
        x = x.flatten(2).transpose(1, 2)

        return x

In other words, the Mlp module effectively performs the following sequence of operations. Since the dropout probability is actually 0, we remove the dropout layers here. For simplicity, bias terms are also omitted.

import torch.nn.functional as F

# H, W = 192, 256
# input = torch.empty(N, H * W, C)

# Step 1: fc1
# fc1_weight = nn.Parameter(torch.empty(hidden, C))
fc1_output = F.linear(input, fc1_weight)

# Step 2: dwconv
N, __, hidden = fc1_output.shape

# Layout conversion: [N, H*W, hidden] -> [N, hidden, H*W] -> [N, hidden, H, W]
fc1_output_transpose = fc1_output.transpose(1, 2).view(N, hidden, H, W)

# dwconv_weight = nn.Parameter(torch.empty([hidden, 1, 3, 3]))
dwconv_output = F.conv2d(
    fc1_output_transpose,
    dwconv_weight,
    stride=1,
    padding=1,
    groups=C_hidden,
)

# Layout conversion: [N, hidden, H, W] -> [N, hidden, H*W] -> [N, H*W, hidden]
dwconv_output_transpose = dwconv_output.flatten(2).transpose(1, 2)

# Step 3: activation
act_output = F.gelu(dwconv_output_transpose)

# Step 4: fc2
# fc2_weight = nn.Parameter(torch.empty(hidden, out_features))
mlp_output = F.linear(act_output, fc2_weight)

The main issue is that the intermediate Conv2d forces layout conversions before and after the convolution. Under fp32, this overhead is not significant. Under int8, however, it accounts for a large share of runtime.

Improvement

To remove the layout-conversion overhead, we convert nn.Linear into nn.Conv2d with a kernel size of 1x1:

class MapMlp(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_features: int,
        out_features: int,
        act_layer: type[nn.Module] = nn.GELU,
        drop: float = 0.0,
    ):
        super().__init__()
        self.fc1 = nn.Conv2d(in_features, hidden_features, 1, bias=True)
        self.dwconv = nn.Conv2d(hidden_features, hidden_features, 3, 1, 1, bias=True, groups=hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=True)
        self.drop = nn.Dropout(drop)

    def forward(self, x: torch.Tensor, height: int, width: int) -> torch.Tensor:
        batch = x.shape[0]
        x = x.reshape(batch, height, width, -1).permute(0, 3, 1, 2).contiguous()
        x = self.fc1(x)
        x = self.dwconv(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x.permute(0, 2, 3, 1).reshape(batch, height * width, -1).contiguous()

Results

The rewritten model performs much better:

The per-operator runtime also changes significantly:

After the rewrite, the runtime of FakeQuantize and Reorder becomes almost negligible.

Summary

After a series of optimizations, CPU throughput improved by nearly 3x.

I had originally planned to further optimize the model with TVM, but the experience was far from plug-and-play. Following the official End-to-End Optimize Model tutorial, the resulting inference latency was 45 seconds. After applying Library Dispatch, with DNNL selected as the library⁵, the latency dropped from 45 seconds to 7.8 seconds. It appears that only conv2d was dispatched, and I did not have the bandwidth to continue digging further.

One final aside: the search experience in the official OpenVINO documentation is, frankly, not very helpful.

The model has been published on Hugging Face: xvyv99/aggpose-lite. The corresponding code is available on GitHub: xvyv99/aggpose-lite.