3、环境准备：Python环境、PyTorch/TensorFlow安装、ONNX与ONNX Runtime安装

3.1 环境准备的核心原理：版本匹配是首要约束

ONNX 模型导出并非单纯的“代码转换”，而是一个跨框架的中间表示（IR）生成过程。其核心原理在于：源框架（PyTorch/TensorFlow）通过各自的 trace 或 graph 机制，将动态计算图静态化为 ONNX 的节点序列。这一过程高度依赖三个层面的版本协同：

  • Python 解释器版本：影响所有二进制扩展（如 torch、tensorflow）的 ABI 兼容性，Python 3.8~3.11 是当前最稳妥区间。
  • 深度学习框架版本：决定了 torch.onnx.export 或 tf2onnx.convert 所支持的算子集（opset）上限与导出 API 行为。
  • ONNX 与 ONNX Runtime 版本：ONNX 负责定义模型结构，ONNX Runtime 负责推理执行。两者版本不匹配时，常出现“导出成功但推理报错”的隐性故障。

下图展示了从源码框架到 ONNX Runtime 推理的完整环境依赖链路，以及各组件之间的版本约束关系：

Python 3.8 ~ 3.11 pip / conda 虚拟环境隔离 PyTorch 2.x torch.onnx.export() opset_version=12~18 TensorFlow 2.x tf2onnx.convert() opset 默认 13+ ONNX 模型文件 (.onnx) 中间表示：算子图 + 权重 + 元信息 ONNX Runtime (CPU / CUDA / TensorRT) 版本需 ≥ ONNX 导出时的 opset 支持范围 版本约束检查点 ① torch 与 opset 上限 ② tf 与 onnx 算子映射 ③ onnxruntime 与 onnx
关键认知：ONNX 文件本身是“静态”的，但它的生成与执行都依赖“动态”的运行时环境。很多导出报错（如 Unsupported operator、Constant folding failed）本质上都是版本链断裂导致的。因此，环境准备不是一次性工作，而是每次新建项目时的首要检查项。

3.2 方法流程：从零搭建可复现的 ONNX 导出环境

推荐使用 conda 或 venv 创建独立虚拟环境，避免污染系统级 Python。以下流程以 PyTorch 路线为主，TensorFlow 路线在步骤 3 中给出对应命令。

步骤 1：创建并激活虚拟环境

# 使用 conda（推荐，便于管理 CUDA 依赖）
conda create -n onnx_env python=3.10 -y
conda activate onnx_env

# 或使用 venv（轻量级）
python -m venv onnx_env
# Windows: onnx_env\Scripts\activate
# Linux/macOS: source onnx_env/bin/activate

步骤 2：安装深度学习框架

# PyTorch 路线（以 CUDA 11.8 为例）
pip install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu118

# TensorFlow 路线（CPU 版）
pip install tensorflow==2.15.0

# TensorFlow 路线（GPU 版，需匹配 CUDA 12.2）
pip install tensorflow[and-cuda]==2.15.0

步骤 3：安装 ONNX 生态核心库

# ONNX 标准库（负责模型定义与检查）
pip install onnx==1.15.0

# ONNX Runtime（负责推理执行）
pip install onnxruntime==1.17.1          # CPU 版
pip install onnxruntime-gpu==1.17.1      # GPU 版（需 CUDA 12.x）

# 辅助工具（强烈建议安装）
pip install onnxconverter-common          # 版本兼容工具
pip install onnxruntime-extensions       # 自定义算子支持

步骤 4：验证安装完整性

python -c "
import torch, onnx, onnxruntime as ort
print('PyTorch:', torch.__version__)
print('ONNX:', onnx.__version__)
print('ONNX Runtime:', ort.__version__)
print('CUDA available:', torch.cuda.is_available())
print('ONNX Runtime providers:', ort.get_available_providers())
"

输出示例（正常情况）：

PyTorch: 2.1.2+cu118
ONNX: 1.15.0
ONNX Runtime: 1.17.1
CUDA available: True
ONNX Runtime providers: ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']

3.3 案例/场景：典型环境组合与验证清单

以下表格列出了三种最常见的生产环境组合，以及对应的验证要点。请根据你的部署目标（CPU 服务器 / GPU 推理 / 边缘设备）选择其中一种作为基线。

场景 Python 深度学习框架 ONNX ONNX Runtime 验证要点
CPU 服务器推理 3.10 PyTorch 2.1.2 1.15.0 onnxruntime 1.17.1 确认 get_available_providers() 仅含 CPU
GPU 高性能推理 3.10 PyTorch 2.1.2+cu118 1.15.0 onnxruntime-gpu 1.17.1 CUDA 版本与 cuDNN 匹配；torch.cuda.is_available() 为 True
TensorFlow 存量模型迁移 3.9 TensorFlow 2.15.0 1.14.1 onnxruntime 1.16.3 tf2onnx 转换时指定 --opset 13 避免算子缺失

实操验证脚本：在完成上述安装后，运行以下代码生成一个最小 ONNX 模型并执行推理，以确认整条链路通畅：

import torch
import onnx
import onnxruntime as ort

# 1. 定义一个最小模型
class TinyModel(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x + 1.0)

model = TinyModel()
dummy_input = torch.tensor([[1.0, -2.0]])

# 2. 导出 ONNX
torch.onnx.export(model, dummy_input, "tiny.onnx",
                  input_names=["input"], output_names=["output"],
                  opset_version=14)

# 3. 验证 ONNX 文件完整性
onnx_model = onnx.load("tiny.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型结构检查通过")

# 4. ONNX Runtime 推理
sess = ort.InferenceSession("tiny.onnx")
result = sess.run(None, {"input": dummy_input.numpy()})
print("推理结果:", result[0])  # 期望输出 [[2.0, 0.0]]

3.4 误区/复盘：环境准备中的高频踩坑点

根据大量导出报错案例的统计，以下四个环境相关问题占据了约 70% 的“非模型本身”故障。请务必在开始模型导出前逐项排查。

误区一：盲目使用最新版本

最新版 PyTorch（如 2.3+）可能默认使用更新的算子实现，导致导出的 ONNX 模型需要更高 opset 才能表达。而 ONNX Runtime 的稳定版往往滞后于最新 opset。建议：PyTorch 固定使用 2.1.x 或 2.2.x，ONNX Runtime 使用 1.17.x 或 1.18.x，这是当前兼容性最成熟的组合。

误区二：忽略 CUDA 版本的三方匹配

PyTorch 的 CUDA 版本、ONNX Runtime GPU 版的 CUDA 版本、以及系统驱动的 CUDA 版本必须形成“向下兼容”链。例如：系统驱动支持 CUDA 12.2，PyTorch 使用 cu118 编译版本，ONNX Runtime 使用 CUDA 12.x 版本——这种混搭会导致 CUDA error: no kernel image available。建议使用 nvidia-smi 查看驱动支持的 CUDA 版本，然后统一选择对应编译版本。

误区三：在系统 Python 中直接安装

系统 Python（如 Ubuntu 自带的 3.8）通常被多个系统工具依赖。直接 pip install 可能触发 externally-managed-environment 错误，或意外升级关键库导致系统故障。务必使用虚拟环境，并在项目根目录保存 requirements.txt 以便复现。

误区四：ONNX 与 ONNX Runtime 版本脱节

ONNX 1.15 支持 opset 最高到 20，但 ONNX Runtime 1.17 仅完整支持 opset 19。若导出时指定 opset_version=20，虽然 ONNX 库能生成文件，但 Runtime 加载时会报 Unsupported opset version。建议：导出前先查询 onnxruntime.__version__ 对应的最大 opset，或统一使用 opset 14~17 这个“安全区间”。

3.5 环境验证清单（快速自查）

在开始任何模型导出工作前，请逐项确认以下清单，确保环境处于健康状态：

  1. python --version 输出为 3.8~3.11 之间的版本。
  2. pip list | grep -E "torch|tensorflow" 确认框架版本与预期一致。
  3. python -c "import onnx; print(onnx.__version__)" 确认 ONNX 版本 ≥ 1.14。
  4. python -c "import onnxruntime; print(onnxruntime.__version__)" 确认 Runtime 版本与 ONNX 相差不超过 2 个 minor 版本。
  5. 运行 3.3 节中的“最小模型验证脚本”，确认导出与推理均无异常。
  6. 若使用 GPU，额外运行 python -c "import torch; print(torch.cuda.is_available())" 确认输出为 True。
← 上一章 📖 返回目录 下一章 →