2. 环境搭建：安装ONNX、ONNX Runtime、PyTorch/TensorFlow

做嵌入式AI部署，第一步就是搭环境。很多人觉得这步简单，直接pip install就完事了。但我在项目里踩过不少坑——版本不兼容、编译报错、跑起来发现CPU不支持某个算子……嗯，这一章咱们就把环境彻底搞定。

2.1 整体环境概览

先看看我们要装哪些东西，它们之间的关系是什么样的。

说白了，整个流程就是：训练框架 → ONNX → ONNX Runtime。我习惯把ONNX比作“通用翻译器”，不管你用PyTorch还是TensorFlow写的模型，都能转成同一种格式，然后在ONNX Runtime上跑。

2.2 安装PyTorch

我个人更常用PyTorch做模型训练，因为它的动态图机制调试起来方便。当然，如果你习惯TensorFlow，完全没问题，后面章节两种框架的转换我都会讲。

2.2.1 确认Python版本

先看看你的Python版本。我建议用Python 3.8~3.11，太老或太新的版本可能会有兼容性问题。

python3 --version
# 输出示例：Python 3.10.12

2.2.2 安装PyTorch

去PyTorch官网（pytorch.org）选对应的配置。这里给一个CPU版本的安装命令——嵌入式端部署一般用CPU推理，GPU版本在开发机上用就行。

# CPU版本（推荐用于嵌入式开发）
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

# 如果你有NVIDIA显卡，可以装CUDA版本
# pip3 install torch torchvision torchaudio

2.3 安装TensorFlow（可选）

如果你只用PyTorch，这节可以跳过。但很多工业项目还是用TensorFlow的，尤其是那些老项目。我遇到过客户给的模型是TensorFlow 1.x的，转ONNX时折腾了好久。

# 安装TensorFlow CPU版本
pip3 install tensorflow

# 验证安装
python3 -c "import tensorflow as tf; print(tf.__version__)"

2.4 安装ONNX

ONNX本身是一个Python包，负责模型导出和格式转换。安装很简单：

pip3 install onnx

# 验证安装
python3 -c "import onnx; print(onnx.__version__)"

这里有个坑——ONNX的版本要和你的训练框架版本匹配。我整理了一个对照表：

2.5 安装ONNX Runtime

ONNX Runtime是推理引擎，负责在嵌入式Linux上加载并运行ONNX模型。它有几个版本：

• onnxruntime — CPU版本，通用
• onnxruntime-gpu — GPU版本，需要CUDA
• onnxruntime-openvino — 针对Intel芯片优化
• onnxruntime-tensorrt — 针对NVIDIA Jetson优化

嵌入式端部署，我一般用CPU版本。安装命令：

# CPU版本（嵌入式Linux推荐）
pip3 install onnxruntime

# 验证安装
python3 -c "import onnxruntime; print(onnxruntime.__version__)"

2.6 验证环境

所有东西装完后，跑一个完整的验证脚本。这个脚本会：创建简单模型 → 转ONNX → 用ONNX Runtime推理 → 对比结果。

import torch
import torch.nn as nn
import onnx
import onnxruntime
import numpy as np

# 1. 创建一个简单的PyTorch模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)
    
    def forward(self, x):
        return self.fc(x)

model = SimpleModel()
model.eval()

# 2. 导出ONNX
dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    "simple_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

# 3. 验证ONNX模型
onnx_model = onnx.load("simple_model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型验证通过！")

# 4. 用ONNX Runtime推理
ort_session = onnxruntime.InferenceSession("simple_model.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)

# 5. 对比PyTorch和ONNX Runtime的结果
with torch.no_grad():
    torch_output = model(dummy_input)

print("PyTorch输出:", torch_output.numpy())
print("ONNX Runtime输出:", ort_outputs[0])
print("最大误差:", np.max(np.abs(torch_output.numpy() - ort_outputs[0])))

如果最后输出的最大误差在1e-5以内，说明环境搭建成功。我见过有人误差很大，结果发现是模型没设成eval模式——训练时的Dropout和BatchNorm在推理时会有差异。

2.7 常见问题与解决

嗯，环境搭建就到这里。记住一个原则：版本对齐。训练框架、ONNX、ONNX Runtime三者的版本要匹配，否则后面会出各种奇怪的问题。我见过太多人因为版本不对，浪费一整天在调试上。