3、环境准备：搭建可复现的调试环境（Python + PyTorch + ONNX Runtime）

做算子融合的精度排查，最怕什么？

我最怕的是——环境不一致。你在本地跑得好好的，放到服务器上就崩了。或者昨天还能复现的bug，今天莫名其妙消失了。说白了，精度问题本身就像个幽灵，如果连环境都不固定，你根本没法跟它斗。

所以这一章，我带你搭一个「铁板一块」的调试环境。我自己的习惯是：用Python 3.8 + PyTorch 1.13 + ONNX Runtime 1.14，这个组合我踩过无数坑，相对稳定。

3.1 为什么必须固定版本？

你想想看，ONNX Runtime每个版本对算子的实现都有微调。PyTorch导出onnx时，不同版本对某些算子的拆解方式也不一样。我曾经遇到过一个问题：同一个模型，PyTorch 1.12导出的onnx里，Gemm被拆成了MatMul+Add；而PyTorch 1.13里，它保留了Gemm。这一拆一合，融合规则就全变了。

所以，版本锁定是第一要务。

3.2 创建干净的虚拟环境

我个人习惯用conda，因为它对CUDA版本的管理比较省心。当然，venv也行，看你喜好。

# 创建环境
conda create -n fusion_debug python=3.8 -y
conda activate fusion_debug

# 安装PyTorch（以CUDA 11.7为例）
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

# 安装ONNX Runtime
pip install onnxruntime-gpu==1.14.1

# 辅助工具
pip install onnx onnxsim netron

嗯，这里要注意：onnxruntime-gpu 和 onnxruntime 是两回事。如果你只是做CPU推理，用后者就行。但做算子融合调试，我建议直接用GPU版本，因为很多融合优化在GPU后端才生效。

3.3 验证环境是否正常

装完了别急着跑模型。先写个最小验证脚本，确保每个组件都能正常工作。

import torch
import onnxruntime as ort
import onnx

print("PyTorch version:", torch.__version__)
print("ONNX Runtime version:", ort.__version__)
print("ONNX version:", onnx.__version__)

# 验证GPU可用
print("CUDA available:", torch.cuda.is_available())

# 验证ONNX Runtime GPU可用
providers = ort.get_available_providers()
print("Available providers:", providers)
assert 'CUDAExecutionProvider' in providers, "GPU provider not found!"

如果输出里没有报错，并且能看到 CUDAExecutionProvider，那环境基本就稳了。

3.4 准备一个「最小复现模板」

做精度排查，最忌讳的是拿整个大模型直接跑。你想想看，一个几百层的模型，融合了十几个算子，你根本不知道是哪一步出了问题。

我的做法是：先构造一个极简的融合场景。比如，一个Conv+BN+ReLU的序列，这是最常见的融合模式。

import torch
import torch.nn as nn

class SimpleFusionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

model = SimpleFusionModel().eval()

# 导出onnx
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, 
    dummy_input,
    "fusion_test.onnx",
    opset_version=17,
    input_names=['input'],
    output_names=['output']
)

print("ONNX model exported successfully.")

这个模板的好处是：你知道它一定会触发Conv+BN+ReLU的融合。如果这个简单case都出问题，那就不用查大模型了，直接定位到融合逻辑本身。

3.5 对比工具：PyTorch vs ONNX Runtime

精度排查的核心，就是对比。对比融合前和融合后的输出差异。我一般准备两个推理函数：

def run_pytorch(model, input_tensor):
    """用PyTorch做推理，作为基准"""
    with torch.no_grad():
        output = model(input_tensor)
    return output.numpy()

def run_onnx(onnx_path, input_tensor):
    """用ONNX Runtime做推理，看融合后的结果"""
    session = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider'])
    input_name = session.get_inputs()[0].name
    output = session.run(None, {input_name: input_tensor.numpy()})
    return output[0]

# 对比
pytorch_out = run_pytorch(model, dummy_input)
onnx_out = run_onnx("fusion_test.onnx", dummy_input)

diff = np.abs(pytorch_out - onnx_out).max()
print(f"Max difference: {diff:.6f}")

如果diff在1e-5以内，说明融合没引入精度问题。如果差了好几个数量级，那就要深入查了。

3.6 环境快照与复现

环境搭好了，怎么保证别人也能复现？我建议你生成一个 requirements.txt，并且把关键包的版本写死。

pip freeze | grep -E "torch|onnx|numpy" > requirements_fusion.txt

然后把这个文件连同你的复现代码一起提交到git仓库。这样，不管是谁，只要执行 pip install -r requirements_fusion.txt，就能得到一模一样的环境。

3.7 本章小结

环境准备这事儿，看着简单，但往往是精度排查的第一道坎。我见过太多人花了两天查bug，最后发现是ONNX Runtime版本不对。所以，别嫌麻烦，把环境锁死，把模板写好，后面才能专心对付真正的精度问题。

想要更多资料，联系可可