Mini-Infer (12): 特征提取的收缩 — `Pooling` 算子与架构复用

引言：不仅仅是卷积

在卷积神经网络（CNN）中，如果说 Conv2D 是“提取特征”的画家，那么 Pooling（池化）就是“提炼精华”的编辑。

没有池化层，特征图（Feature Map）的尺寸会一直保持不变（或仅缓慢减小），这将导致计算量爆炸，且网络难以学习到具有“平移不变性”的高层特征。

1. 定义池化：`PoolingParam` 与 TensorRT 对齐

struct PoolingParam : public OpParam {
    PoolingType type;       // MAX or AVERAGE
    int kernel_h, kernel_w;
    int stride_h, stride_w;
    int padding_h, padding_w;
    // ...
};

这里有两个值得注意的设计选择：

支持非对称参数：kernel_h vs kernel_w，padding_h vs padding_w。很多简单的框架只支持正方形核，但支持非对称是处理 NLP 任务或特殊长宽比图片的关键。
TensorRT 对齐：明确注释了参考 TensorRT 的 IPoolingLayer。特别是对于 AVERAGE 池化，采用了 “Exclude Padding” (计算平均值时不把 Padding 的 0 算入分母) 的策略。这是推理引擎中防止边缘特征被稀释的标准做法。

2. 算子实现：`Pooling` Operator

Pooling 算子的实现逻辑（operators/pooling.cpp）展示了标准的算子开发流程：

A. 参数校验 (Validation)

算子不仅仅是计算，更是“守门员”。在 Pooling::Pooling 构造函数和 forward 开头，你做了详尽的检查：

核大小必须 > 0
步长必须 > 0
Padding 必须 >= 0
输入维度必须是 4 (NCHW)

这些检查在调试模型转换（如 ONNX Parser）时会节省大量时间。

B. 形状推导 (`infer_shape`)

这是 Engine 进行静态内存规划的依据。池化的输出尺寸计算公式是经典的：

$H_{out} = \lfloor \frac{H_{in} + 2 \times padding - kernel}{stride} \rfloor + 1$

1 2	int H_out = (H_in + 2 * param_.padding_h - param_.kernel_h) / param_.stride_h + 1; int W_out = (W_in + 2 * param_.padding_w - param_.kernel_w) / param_.stride_w + 1;

C. 计算分发 (`forward`)

Pooling::forward 再次展示了我们 KernelRegistry 的威力。算子本身不包含任何 for 循环，它只是负责：

解析参数。
准备输入输出指针。
分发 (Dispatch) 给对应的 Kernel。

// 极其干净的分发逻辑
if (param_.type == PoolingType::MAX) {
    kernels::PoolingKernel::maxpool2d<float>(...);
} else if (param_.type == PoolingType::AVERAGE) {
    kernels::PoolingKernel::avgpool2d<float>(...);
}

3. 内核实现：`PoolingKernel` (CPU)

虽然池化比卷积简单，但在实现细节上很容易出错。你的 CPU 实现 (kernels/pooling.cpp) 处理得非常严谨。

Max Pooling 的细节

MaxPool 的核心是在窗口中找最大值。

陷阱：如果初始化为 0，当输入特征包含负数（例如 LeakyReLU 的输出）时，结果会出错。
你的实现：正确使用了 std::numeric_limits<T>::lowest()（即 -FLT_MAX）作为初始值。

1
2
3

T max_val = std::numeric_limits<T>::lowest();
// ... 循环 ...
max_val = std::max(max_val, val);

Average Pooling 的细节

AvgPool 的核心是处理 Padding。

策略：TensorRT 风格的 count_include_pad=false。
你的实现：只累加有效范围内的像素，并用 count 变量记录实际参与计算的像素个数，而不是简单地除以 kernel_h * kernel_w。

// 你的代码：Exclude Padding 的正确实现
int count = 0;
for (int h = h_start; h < h_end; ++h) {
    // ...
    sum += input_nc[...];
    ++count; // 只统计有效像素
}
output_nc[...] = (count > 0) ? (sum / count) : 0;

4. 架构复用的胜利

实现 Pooling 算子最大的感触应该是：“这也太快了吧？”

我们不需要重写 KernelRegistry 模板，直接 DEFINE_REGISTRY_ALIAS 即可。
我们不需要重写 Buffer 或 Tensor 管理。
我们不需要重写后端分发逻辑。

我们只专注于 Pooling 特有的逻辑（滑动窗口公式）。这就是一个成熟架构带来的红利——添加新算子的边际成本越来越低。

总结

通过添加 Pooling 算子，Mini-Infer 已经集齐了构建 VGG 网络的所有拼图：

Conv2D: 特征提取
ReLU: 非线性激活
Pooling: 降采样与特征压缩
Linear: 分类