📚 课程目录第2章环境搭建

2、环境搭建：安装PyTorch、ONNX、onnxruntime，配置开发环境

在开始ONNX模型导出实战之前，必须先搭建一套干净、可复现的开发环境。很多初学者在后续导出过程中遇到的“算子不支持”“版本不兼容”“推理结果不一致”等问题，根源往往不在代码逻辑，而在于环境组件版本错配。本节将带你完成PyTorch、ONNX、onnxruntime的安装与验证，并给出推荐的版本组合与配置策略。

2.1 核心组件与版本匹配原理

ONNX模型导出链路涉及三个关键组件，它们各自扮演不同角色：

PyTorch（torch）：模型定义与训练框架，负责将动态计算图转换为静态ONNX图。其内部通过 torch.onnx.export 调用TorchScript追踪器完成导出。
ONNX（onnx）：中间表示格式的标准库，提供模型结构校验、节点遍历、版本转换等能力。它不参与推理，但负责保证模型文件的规范性与跨框架兼容性。
onnxruntime（ort）：跨平台推理引擎，直接加载ONNX模型并执行前向计算。它是验证导出结果正确性的“裁判”，也是后续部署到生产环境的核心依赖。

三者之间存在严格的版本联动关系：PyTorch导出ONNX的算子集版本（opset）必须被ONNX库支持，同时onnxruntime的算子实现也必须覆盖该opset。如果PyTorch版本过新而onnxruntime版本过旧，很可能出现“导出成功但推理报错”的尴尬局面。

推荐版本组合（截至2025年初稳定方案）：Python 3.10 + PyTorch 2.3.x + onnx 1.16.x + onnxruntime 1.18.x。该组合支持opset 17~19，覆盖绝大多数CV与NLP模型的导出需求。

2.2 安装方法流程

推荐使用 conda 或 venv 创建独立虚拟环境，避免污染系统Python。以下以conda为例，展示完整安装流程：

# 1. 创建虚拟环境（Python 3.10）
conda create -n onnx_env python=3.10 -y
conda activate onnx_env

# 2. 安装PyTorch（CPU版示例，GPU版请根据CUDA版本调整）
pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cpu

# 3. 安装ONNX与onnxruntime
pip install onnx==1.16.2 onnxruntime==1.18.1

# 4. 安装辅助工具（可选但推荐）
pip install onnxscript onnxruntime-extensions netron

  GPU环境注意事项：若需GPU加速，请先运行 nvidia-smi 确认CUDA版本，然后访问 https://pytorch.org/get-started/locally/ 获取匹配的安装命令。例如CUDA 12.1对应：pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/cu121。onnxruntime的GPU版本需单独安装：pip install onnxruntime-gpu==1.18.1。

2.3 环境验证与常见问题排查

安装完成后，必须执行一套完整的验证脚本，确认三个组件能协同工作。以下代码将创建一个简单的线性模型，导出为ONNX，再用onnxruntime加载并对比输出：

import torch
import onnx
import onnxruntime as ort
import numpy as np

# 1. 定义并导出模型
class LinearModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(4, 2)

    def forward(self, x):
        return self.fc(x)

model = LinearModel().eval()
dummy_input = torch.randn(1, 4)

torch.onnx.export(
    model, dummy_input, "linear.onnx",
    input_names=["input"], output_names=["output"],
    opset_version=17, dynamic_axes={"input": {0: "batch"}}
)

# 2. 校验ONNX模型结构
onnx_model = onnx.load("linear.onnx")
onnx.checker.check_model(onnx_model)
print(f"ONNX模型版本: {onnx_model.ir_version}")

# 3. onnxruntime推理对比
ort_session = ort.InferenceSession("linear.onnx")
ort_input = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_output = ort_session.run(None, ort_input)[0]

# 4. 对比PyTorch与ONNX输出
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

print("PyTorch输出:", torch_output)
print("ONNX输出:  ", ort_output)
print("最大误差:", np.abs(torch_output - ort_output).max())

预期输出中，最大误差应小于 1e-5。若误差过大，说明导出过程中存在数值精度问题，需检查模型是否包含不稳定的操作（如某些归一化层）。

验证项	命令/方法	预期结果
PyTorch版本	`python -c "import torch; print(torch.__version__)"`	2.3.x
ONNX版本	`python -c "import onnx; print(onnx.__version__)"`	1.16.x
onnxruntime版本	`python -c "import onnxruntime; print(onnxruntime.__version__)"`	1.18.x
模型结构校验	`onnx.checker.check_model(onnx_model)`	无异常抛出
推理一致性	PyTorch输出 vs ONNX输出	最大误差 < 1e-5

2.4 典型场景：从零搭建可复现的导出环境

假设你正在为一个图像分类项目准备ONNX导出环境，项目需要支持动态批量大小和多种输入分辨率。此时环境配置需要额外考虑以下因素：

固定随机种子：在导出前设置 torch.manual_seed(42)，确保模型权重初始化可复现。
设置推理模式：调用 model.eval() 并包裹 torch.no_grad()，禁用Dropout和BatchNorm的训练行为。
动态轴配置：在 torch.onnx.export 中通过 dynamic_axes 参数声明动态维度，避免模型被固定为静态shape。
算子集版本选择：优先使用 opset_version=17，该版本在算子覆盖率和兼容性之间取得较好平衡。

# 实际项目中的导出配置示例
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    input_names=["images"],
    output_names=["logits"],
    dynamic_axes={
        "images": {0: "batch", 2: "height", 3: "width"},
        "logits": {0: "batch"}
    },
    opset_version=17,
    do_constant_folding=True,  # 启用常量折叠优化
    export_params=True         # 导出模型权重
)

常见误区警示：不要直接在训练脚本中执行导出操作。训练过程中模型处于 train() 模式，BatchNorm层会使用批内统计量而非全局统计量，导致导出的ONNX模型推理结果异常。务必先加载训练好的权重，切换到 eval() 模式后再导出。

2.5 环境配置复盘与优化建议

完成基础安装后，建议进行以下三项优化，提升后续开发效率：

使用requirements.txt锁定版本：将精确版本号写入文件，便于团队协作和CI/CD流水线复现。
配置Netron可视化工具：安装Netron后，可在浏览器中直接打开ONNX模型文件，直观查看网络结构、节点参数和输入输出shape，极大方便调试。
建立模型导出测试集：准备3~5个不同结构的模型（CNN、Transformer、MLP），每次修改环境后先运行测试集，快速发现版本兼容性问题。

# requirements.txt 示例
torch==2.3.1
torchvision==0.18.1
onnx==1.16.2
onnxruntime==1.18.1
onnxscript==0.1.2
netron==7.6.0

最后，验证环境是否支持GPU推理（如果安装了GPU版本）：

import onnxruntime as ort
print("可用执行提供程序:", ort.get_available_providers())
# 预期输出包含 'CUDAExecutionProvider'（GPU版）或 'CPUExecutionProvider'