2. 环境准备：安装ONNX、ONNX Runtime、PyTorch/TensorFlow

说实话，环境配置这件事，看着简单，但坑是真不少。我见过太多同学花了一整天装环境，最后发现版本对不上，或者装完跑不起来。这一章，咱们就把这事一次性搞定。

我会带着你一步步装好ONNX、ONNX Runtime，还有PyTorch和TensorFlow。最后再跑个验证脚本，确保所有东西都能正常干活。

2.1 整体架构：这些工具到底怎么配合？

先别急着装东西。咱们得搞清楚这些组件之间的关系。我画了张图，你看一眼就明白了。

说白了，流程就三步：

• PyTorch/TensorFlow 负责训练模型
• ONNX 作为中间桥梁，把模型转成统一格式
• ONNX Runtime 负责加载并加速推理

嗯，这里要注意：ONNX 本身不是推理引擎，它只是个模型描述格式。真正干活的是 ONNX Runtime。

2.2 安装Python环境

我个人习惯用 conda 管理环境。为什么？因为项目多了，版本冲突是家常便饭。conda 能帮你隔离得干干净净。

创建环境：

conda create -n onnx_env python=3.9 -y
conda activate onnx_env

Python 版本选 3.8-3.10 都行。我个人推荐 3.9，兼容性最好，踩坑最少。

2.3 安装ONNX核心库

ONNX 的安装其实很简单。但版本选择有讲究。

pip install onnx==1.15.0

为什么指定版本？因为 ONNX 的 op 集版本和框架版本是绑定的。装太新的 ONNX，老框架可能导不出模型；装太旧的，新算子又不支持。

2.4 安装ONNX Runtime

ONNX Runtime 分好几个版本。你想想看，你的模型最终跑在什么设备上？

• CPU 推理：装 onnxruntime
• GPU 推理：装 onnxruntime-gpu
• 边缘设备：装 onnxruntime-tensorrt 或 onnxruntime-openvino

我建议你两个都装上，调试时用 CPU，上线时切 GPU。

# CPU 版本
pip install onnxruntime==1.16.3

# GPU 版本（需要 CUDA 11.8 + cuDNN 8.6）
pip install onnxruntime-gpu==1.16.3

2.5 安装PyTorch

PyTorch 的安装，我建议去官网查命令。因为 CUDA 版本不同，命令也不一样。

# CPU 版本
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cpu

# GPU 版本（CUDA 11.8）
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118

# GPU 版本（CUDA 12.1）
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu121

装完后验证一下：

import torch
print(torch.__version__)        # 2.1.0
print(torch.cuda.is_available()) # True 表示 GPU 可用

2.6 安装TensorFlow

TensorFlow 2.x 的安装也类似。不过要注意，TF 2.13 之后对 CUDA 版本要求更严格了。

# CPU 版本
pip install tensorflow==2.13.0

# GPU 版本
pip install tensorflow==2.13.0

# 额外安装 tf2onnx 用于导出
pip install tf2onnx==1.15.1

验证：

import tensorflow as tf
print(tf.__version__)               # 2.13.0
print(tf.config.list_physical_devices('GPU'))  # 能看到 GPU 设备

2.7 完整环境验证

所有东西装完后，跑一个完整的验证脚本。这一步不能省。

import onnx
import onnxruntime
import torch
import tensorflow as tf

print("=" * 50)
print("环境验证报告")
print("=" * 50)

# 1. ONNX 版本
print(f"ONNX 版本: {onnx.__version__}")

# 2. ONNX Runtime 版本和设备
print(f"ONNX Runtime 版本: {onnxruntime.__version__}")
print(f"ONNX Runtime 设备: {onnxruntime.get_device()}")

# 3. PyTorch 版本和 GPU
print(f"PyTorch 版本: {torch.__version__}")
print(f"PyTorch GPU 可用: {torch.cuda.is_available()}")

# 4. TensorFlow 版本和 GPU
print(f"TensorFlow 版本: {tf.__version__}")
print(f"TensorFlow GPU 可用: {len(tf.config.list_physical_devices('GPU')) > 0}")

print("=" * 50)
print("环境验证完成！")
print("=" * 50)

输出示例：

==================================================
环境验证报告
==================================================
ONNX 版本: 1.15.0
ONNX Runtime 版本: 1.16.3
ONNX Runtime 设备: GPU
PyTorch 版本: 2.1.0
PyTorch GPU 可用: True
TensorFlow 版本: 2.13.0
TensorFlow GPU 可用: True
==================================================
环境验证完成！
==================================================

2.8 常见问题速查