Mini-Infer (11): 下采样利器 — `Pooling` 算子与架构复用之美

1. 架构红利：零成本的扩展性

回想我们在 Blog 7.6 中付出的努力——我们重构了内核注册表，引入了模板元编程。现在，回报来了。

我们要添加 MaxPool 和 AvgPool，不需要重写任何注册逻辑。只需要几行宏定义：

// mini_infer/kernels/pooling.h

// 1. 定义函数签名
template<typename T>
using MaxPool2DFunc = void (*)(...);

template<typename T>
using AvgPool2DFunc = void (*)(...);

// 2. 【核心】一键生成注册表！
DEFINE_REGISTRY_ALIAS(MaxPool2DRegistry, MaxPool2DFunc);
DEFINE_REGISTRY_ALIAS(AvgPool2DRegistry, AvgPool2DFunc);

这就完成了！我们瞬间拥有了两个支持自动后端分发（CPU/CUDA）、支持多种数据类型（float/int）的注册表。这就是基础设施的力量。

2. `MaxPool2D`：寻找最强特征

最大池化（Max Pooling）的核心逻辑是在一个滑动窗口内寻找最大值。这看似简单，但有一个经典的陷阱：初始化。

陷阱：如何处理负数输入？

如果我们将最大值的初始值设为 0，当输入全为负数时（虽然在 ReLU 后不常见，但在其他激活函数或中间层中可能出现），输出就会错误的变成 0。

正确的做法： 初始化为该数据类型的**“最小可能值”**（负无穷）。

// mini_infer/kernels/pooling.cpp -> maxpool2d_impl

T max_val = std::numeric_limits<T>::lowest(); // float 是 -FLT_MAX

for (int h = h_start; h < h_end; ++h) {
    for (int w = w_start; w < w_end; ++w) {
        T val = input_nc[h * W_in + w];
        max_val = std::max(max_val, val);
    }
}
output_nc[...] = max_val;

边界处理

我们的实现使用了 TensorRT 风格的逻辑：Padding 区域不参与计算。我们通过 std::max(h_start, 0) 和 std::min(h_end, H_in) 将循环限制在有效的图像区域内。对于 MaxPool，这等价于 Padding 区域是负无穷，不会被选中。

3. `AvgPool2D`：平均值的细节

平均池化（Average Pooling）计算窗口内的平均值。这里的陷阱在于：Padding 算不算分母？

在深度学习框架中，通常有两个选项：

Count Include Pad: 分母总是 Kernel_H * Kernel_W。
Count Exclude Pad: 分母是窗口内实际存在的像素个数。

Mini-Infer（参考 TensorRT 默认行为）选择 Exclude Pad 模式。

// mini_infer/kernels/pooling.cpp -> avgpool2d_impl

T sum = 0;
int count = 0; // 计数器

for (int h = h_start; h < h_end; ++h) {
    for (int w = w_start; w < w_end; ++w) {
        // 只有在有效范围内的像素才累加
        sum += input_nc[h * W_in + w];
        ++count;
    }
}

// 分母使用实际有效的 count，而不是固定的 kernel size
output_nc[...] = (count > 0) ? (sum / count) : 0;

这种处理方式在图像边缘更加自然，避免了因为 Padding 的零值拉低了边缘的平均值。

4. 独立通道处理 (Channel Independence)

与 Conv2D 不同，Pooling 是 Channel-wise 的操作。通道之间没有交互。

这意味着我们的循环结构是：

for (n in Batch) {
    for (c in Channels) { // 每一个通道独立处理
        // ... 对 H * W 进行滑动窗口 ...
    }
}

这种互不依赖的特性，使得 Pooling 算子极易被并行化。在未来的优化中，我们可以轻易地在 n 和 c 维度上使用 OpenMP 多线程，或者在 GPU 上将每个 (n, c) 映射为一个 CUDA Block。

5. 总结

我们在本篇中完成了：

Pooling 接口定义：使用 KernelRegistry 快速定义了 MaxPool 和 AvgPool 的内核接口。
CPU 实现：编写了鲁棒的 C++ 参考实现，正确处理了初始化极值和 Padding 计数问题。
自动注册：在 register_pooling_kernels 中完成了 CPU 后端的注册。