Mini-Infer (13): 端到端验证 — LeNet-5 实战与 PyTorch 对齐

1. 为什么需要端到端测试？

单元测试（Unit Test）只能保证单个算子（如 Conv2D）在特定输入下是正确的。但当几十个算子串联成一个网络时，微小的误差（如 Padding 处理、NCHW vs NHWC 布局差异、float 精度累积）可能会被放大，导致最终分类错误。

端到端测试的目标：

权重加载：验证我们能否正确读取 PyTorch 导出的二进制权重。
计算精度：验证 Mini-Infer 的输出 logits 与 PyTorch 的差异是否在允许范围内（如 1e-5）。
流程打通：验证从图片预处理到最终分类的整个链路。

2. 训练与导出：PyTorch 侧准备 (`lenet5_model.py` & `train_lenet5.py`)

首先，我们需要一个“标准答案”。我们在 PyTorch 中定义并训练一个经典的 LeNet-5。

关键细节：

模型定义：我们严格遵循 Conv -> ReLU -> MaxPool 的顺序，这与我们在 Mini-Infer 中搭建的流程一致。
数据预处理：transforms.Normalize((0.1307,), (0.3081,))。这是 MNIST 的标准归一化。我们在 C++ 推理时必须执行完全相同的预处理，否则结果将毫无意义。

# lenet5_model.py
class LeNet5(nn.Module):
    def __init__(self):
        # ... 定义层 ...
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        # ...

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x))) # 这里的顺序很重要
        # ...

训练脚本 (train_lenet5.py) 会保存 lenet5_best.pth。之后，我们需要使用之前编写的 export_lenet5.py（未展示，但逻辑很简单）将 .pth 转换为我们自定义的二进制权重格式。

3. 推理实现：C++ 侧实战 (`lenet5_inference.cpp`)

这是 Mini-Infer 的“高光时刻”。我们将手动“拼装”出 LeNet-5。

A. 模型组装 (Model Assembly)

不同于 PyTorch 的 nn.Module，在 C++ 中我们需要显式地管理内存和算子。

// lenet5_inference.cpp
class LeNet5 {
public:
    LeNet5(const utils::LeNet5Weights& weights) {
        // 1. 创建算子 (Operators)
        // Conv1: 1->6, 5x5 kernel
        conv1_ = std::make_shared<operators::Conv2D>(
            operators::Conv2DParam(5, 5, 1, 1, 0, 0));
            
        // Pool: 2x2 MaxPool
        pool_ = std::make_shared<operators::Pooling>(
            operators::PoolingParam(operators::PoolingType::MAX, 2, 2, 2, 2));
            
        // ReLU & Linear ...
    }
    // ...

B. 前向传播 (Forward Pass)

因为我们还没有实现 ONNX Parser（这是 B 轨任务），所以目前我们是手动硬编码计算图的执行顺序。这虽然繁琐，但非常有助于理解数据流。

// lenet5_inference.cpp
std::shared_ptr<core::Tensor> forward(std::shared_ptr<core::Tensor> input) {
    auto x = input;
    
    // Layer 1: Conv -> ReLU -> Pool
    conv1_->forward({x, weights_.conv1_weight, weights_.conv1_bias}, outputs);
    x = outputs[0];
    
    relu_->forward({x}, outputs);
    x = outputs[0];
    
    pool_->forward({x}, outputs);
    x = outputs[0];
    
    // ... Layer 2 ...
    // ... Flatten ...
    // ... FC Layers ...
    
    return x;
}

C. 结果验证

我们将 C++ 的推理结果保存为 JSON 文件，其中包含每个样本的 Logits、Probabilities 和预测类别。

4. 自动化测试脚本：`test_lenet5.sh`

为了让验证过程自动化，我们编写了一个 Shell 脚本，它串联了 Python 和 C++：

生成基准：运行 generate_reference_outputs.py，用 PyTorch 跑一遍测试集，保存结果。
运行推理：运行编译好的 lenet5_inference C++ 程序，保存结果。
对比：运行 compare_outputs.py，对比两个 JSON 文件。

# test_lenet5.sh
echo "Step 1: Generating PyTorch Reference Outputs"
python3 generate_reference_outputs.py ...

echo "Step 2: Running C++ Mini-Infer Inference"
../../../build/examples/lenet5_inference ...

echo "Step 3: Comparing Outputs"
python3 compare_outputs.py ...