3. 参考实现搭建：如何搭建基于PyTorch/TensorFlow的参考模型，获取标准输出

好，咱们进入第三个关键环节——参考实现搭建。

说白了，这一步就是给你的算子精度对齐找一个“标准答案”。你想想看，没有标准答案，你怎么知道你的芯片算对了没有？我见过不少团队，一上来就急着写算子映射，结果精度对不齐，回头查半天才发现是参考模型本身就有问题。嗯，这个坑我踩过，所以今天重点聊聊怎么搭一个靠谱的参考模型。

3.1 为什么需要参考模型？

参考模型的作用很简单：提供一组确定性的、可复现的标准输出。你的芯片算子跑出来的结果，要和这个标准输出做对比。误差在允许范围内，就算对齐通过。

我个人习惯把参考模型分成两类：

框架原生模型：直接用PyTorch或TensorFlow的官方API搭建，比如torch.nn.Conv2d、tf.nn.conv2d。这类模型最权威，但灵活性差一些。
自定义参考实现：针对某些特殊算子（比如量化算子、自定义激活函数），框架没有现成的，就得自己写。这时候要格外小心，别把bug带进去。

核心原则：参考模型必须独立于你的芯片实现。不要用同一份代码既跑芯片又跑参考，否则你根本分不清是芯片错了还是参考错了。

3.2 基于PyTorch搭建参考模型

PyTorch是我最常用的参考框架。原因很简单——动态图，调试方便。我在项目中遇到过好几次，用TensorFlow静态图搭参考模型，中间某个节点shape不对，排查起来特别痛苦。换成PyTorch，print大法直接搞定。

3.2.1 基本步骤

固定随机种子：确保每次运行结果一致。
设置确定性算法：关闭cuDNN的自动调优。
构造输入数据：用torch.randn生成随机张量，或者从文件加载真实数据。
前向推理：调用模型得到输出。
保存输出：用torch.save或numpy.save存成文件。

来看一个具体的例子。假设我们要对齐一个Conv2d + ReLU的算子组合：

import torch
import torch.nn as nn
import numpy as np

# 1. 固定随机种子
torch.manual_seed(42)
np.random.seed(42)

# 2. 设置确定性
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# 3. 构造输入 (batch=1, channels=3, height=32, width=32)
input_tensor = torch.randn(1, 3, 32, 32)

# 4. 定义参考模型
class RefModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

model = RefModel()
model.eval()  # 注意：一定要切换到eval模式，关闭dropout/batchnorm

# 5. 前向推理
with torch.no_grad():
    output = model(input_tensor)

# 6. 保存输出
torch.save(output, 'ref_output_conv_relu.pt')
np.save('ref_output_conv_relu.npy', output.numpy())

print(f"输出shape: {output.shape}")
print(f"输出范围: [{output.min().item():.4f}, {output.max().item():.4f}]")

小技巧：保存输出时，我建议同时保存.pt和.npy两种格式。.pt保留PyTorch的元数据，.npy方便C++侧读取。曾经有一次，我忘了保存npy格式，结果芯片侧读不了.pt文件，又得重新跑一遍，浪费时间。

3.3 基于TensorFlow搭建参考模型

如果你的芯片主要对接TensorFlow生态，那参考模型也得用TF搭。不过说实话，TF的调试体验不如PyTorch，所以我一般只在客户明确要求时才用。

3.3.1 关键点

使用tf.function时注意：它会做图优化，可能导致数值精度变化。我建议先用@tf.function(jit_compile=False)关闭XLA编译。
数据类型要一致：TF默认float32，但有些算子内部会用float16累加。记得显式指定dtype=tf.float32。
随机种子：tf.random.set_seed(42)。

同样的Conv2d+ReLU例子，用TF实现：

import tensorflow as tf
import numpy as np

# 固定随机种子
tf.random.set_seed(42)
np.random.seed(42)

# 构造输入
input_tensor = tf.random.normal([1, 32, 32, 3])  # NHWC格式

# 定义参考模型
class RefModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.conv = tf.keras.layers.Conv2D(16, 3, strides=1, padding='same')
        self.relu = tf.keras.layers.ReLU()
    
    def call(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

model = RefModel()
# 注意：TF的Conv2D默认使用channels_last (NHWC)
# 如果你的芯片是NCHW格式，记得做转置

# 前向推理
output = model(input_tensor, training=False)

# 保存
np.save('ref_output_tf_conv_relu.npy', output.numpy())
print(f"输出shape: {output.shape}")

注意：PyTorch默认NCHW，TensorFlow默认NHWC。这是最容易踩的坑。我曾经因为没注意这个，对了一个星期的精度，最后发现是数据排布问题。所以，一定要在保存输出前确认数据格式。

3.4 参考模型的验证与调试

模型搭完了，怎么知道它是对的？我一般做三件事：

数值合理性检查：输出有没有NaN？有没有Inf？范围是否合理？比如ReLU的输出应该全是非负数。
与框架官方文档对比：比如Conv2d的padding行为，PyTorch的'valid'和'same'与TF的略有不同。查文档确认。
小规模手动计算：构造一个1x1的输入，手动算一遍输出，和模型结果对比。这个方法虽然土，但最有效。

举个例子，验证上面的Conv2d：

# 手动验证：输入全1，卷积核全1，bias为0
input_manual = torch.ones(1, 1, 3, 3)
conv_manual = nn.Conv2d(1, 1, kernel_size=2, stride=1, padding=0, bias=False)
# 手动设置卷积核为全1
with torch.no_grad():
    conv_manual.weight.fill_(1.0)

output_manual = conv_manual(input_manual)
# 手动计算：2x2区域求和，结果应该是4
print(f"手动验证输出: {output_manual}")  # 应该输出 [[[[4.]]]]

3.5 保存与组织参考输出

我建议建立一个清晰的目录结构，方便后续对齐时查找：

ref_outputs/
├── conv2d_relu/
│   ├── input.npy
│   ├── output.npy
│   └── config.json      # 记录算子参数：kernel_size, stride, padding等
├── matmul/
│   ├── input_a.npy
│   ├── input_b.npy
│   ├── output.npy
│   └── config.json
└── ...

config.json里记录什么？我一般写这些：

框架版本（PyTorch 1.12.0 / TF 2.10.0）
算子参数（kernel_size, stride, padding, dilation等）
输入shape和dtype
随机种子
生成时间

为什么要记这些？因为框架版本升级可能导致数值精度变化。有一次我升级了PyTorch小版本，结果某个算子的输出变了0.001%，排查了半天才发现是框架内部实现改了。有了config.json，就能快速定位问题。

3.6 常见问题与避坑

我曾经踩过的坑：

忘记设置model.eval()：训练模式和推理模式的BatchNorm、Dropout行为不同。我有一回没设eval，结果每次跑出来的输出都不一样，还以为是随机种子没固定。
数据类型隐式转换：PyTorch的nn.Conv2d默认float32，但如果你输入是float16，它会自动转换。这可能导致精度差异。建议显式统一dtype。
cuDNN的非确定性算法：某些卷积实现（如torch.backends.cudnn.deterministic=False）在GPU上每次结果略有不同。一定要设为True。
多卡推理：DataParallel或DistributedDataParallel会改变数值分布。参考模型建议单卡运行。

3.7 本章小结

搭建参考模型，说白了就是给你的芯片算子找一个“靠谱的标杆”。

核心要点就三条：

固定一切随机因素：种子、确定性算法、eval模式。
数据格式要统一：NCHW vs NHWC，float32 vs float16，提前约定好。
验证参考模型本身：手动计算、文档对比、数值合理性检查，一个都不能少。

嗯，参考模型搭好了，下一步就是拿着这个标准输出去和芯片结果做对比了。不过那是下一章的内容，咱们先把这一步做扎实。

个人建议：如果你时间充裕，可以用PyTorch和TensorFlow各搭一份参考模型，交叉验证。两份结果一致，那基本就稳了。我在一个量产项目中就是这么做的，后来芯片一次流片成功，精度对齐只花了三天。

📚 课程目录

3. 参考实现搭建：如何搭建基于PyTorch/TensorFlow的参考模型，获取标准输出

3.1 为什么需要参考模型？

3.2 基于PyTorch搭建参考模型

3.2.1 基本步骤

3.3 基于TensorFlow搭建参考模型

3.3.1 关键点

3.4 参考模型的验证与调试

3.5 保存与组织参考输出

3.6 常见问题与避坑

3.7 本章小结