2. 环境搭建与工具链准备

说实话,做模型压缩和特征迁移,最怕什么?不是算法调不通,而是环境先崩了。我见过太多同学花两周调通了模型,结果转 ONNX 时卡在算子兼容性上,一查发现是 CUDA 版本没对齐。嗯,这种坑我踩过不止一次。

所以这一章,咱们先把地基打牢。我会带你走一遍 PyTorch、TensorFlow、ONNX、TensorRT、NCNN、OpenVINO 的安装流程,顺便把验证方法也交代清楚。你跟着做一遍,后面 28 章基本不会因为环境问题翻车。

2.1 深度学习框架安装

我个人习惯用 PyTorch 做研究原型,用 TensorFlow 做生产部署。但这不是绝对的,你完全可以根据团队技术栈选一个主攻。下面我把两个框架的安装要点都列出来。

2.1.1 PyTorch 环境配置

PyTorch 的安装其实就一行命令,但坑在 CUDA 版本匹配上。我建议你先用 nvidia-smi 看一下驱动支持的 CUDA 最高版本,然后去 PyTorch 官网选对应的安装命令。

我的习惯: 用 conda 创建独立环境,避免把系统 Python 搞乱。每个项目一个环境,互不干扰。
# 创建 Python 3.9 环境
conda create -n deploy_opt python=3.9 -y
conda activate deploy_opt

# 安装 PyTorch 2.1.0(CUDA 11.8)
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

# 验证安装
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

验证输出应该是类似 2.1.0+cu118 和 True。如果 CUDA 不可用,先检查驱动,别急着重装 PyTorch。

我曾经踩过的坑: 在服务器上装了 PyTorch 1.12,结果同事的代码依赖 torchvision 0.13 的新特性,两个版本不兼容。后来我养成了习惯——每次建环境都把 PyTorch、torchvision、torchaudio 版本写进 requirements.txt,方便复现。

2.1.2 TensorFlow 环境配置

TensorFlow 2.x 的安装比 1.x 友好多了。但要注意,TensorFlow 对 CUDA 和 cuDNN 的版本要求非常严格,差一个小版本都可能报错。

# 安装 TensorFlow 2.13.0
pip install tensorflow==2.13.0

# 验证安装
python -c "import tensorflow as tf; print(tf.__version__); print(tf.config.list_physical_devices('GPU'))"

如果 GPU 列表为空,大概率是 CUDA/cuDNN 版本没对齐。我建议你直接查 TensorFlow 官方版本对应表,别自己猜。

框架 推荐版本 CUDA 版本 cuDNN 版本 Python 版本
PyTorch 2.1.0 11.8 8.7 3.8-3.11
TensorFlow 2.13.0 11.8 8.6 3.8-3.11

2.2 中间表示层:ONNX 安装与验证

ONNX 是模型转换的「通用语言」。不管你是 PyTorch 还是 TensorFlow,最终都要转成 ONNX 才能对接 TensorRT、NCNN 这些推理引擎。

# 安装 ONNX 核心库和 PyTorch 导出工具
pip install onnx onnxruntime onnxruntime-gpu

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

验证通过后,我建议你跑一个简单的导出测试:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.relu(self.conv(x))

model = SimpleModel()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "test.onnx", 
                  input_names=["input"], output_names=["output"],
                  dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}})

# 验证 ONNX 模型
import onnx
onnx_model = onnx.load("test.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型验证通过!")
关键点: dynamic_axes 参数一定要加。不加的话,导出的 ONNX 模型只能接受固定 batch size,部署时非常不灵活。我在项目中吃过这个亏,后来每次导出都默认加上。

2.3 推理引擎安装

推理引擎是模型落地的最后一公里。不同硬件平台选不同的引擎:NVIDIA GPU 用 TensorRT,ARM 设备用 NCNN,Intel CPU 用 OpenVINO。下面我逐个讲安装要点。

2.3.1 TensorRT 安装

TensorRT 的安装稍微麻烦一点,因为它需要和 CUDA 版本严格对应。我建议你直接从 NVIDIA 官网下载 tar 包手动安装,别用 pip,因为 pip 版本经常滞后。

# 下载 TensorRT 8.6.1(对应 CUDA 11.8)
# 解压后设置环境变量
export TRT_PATH=/path/to/TensorRT-8.6.1.6
export LD_LIBRARY_PATH=$TRT_PATH/lib:$LD_LIBRARY_PATH
export PATH=$TRT_PATH/bin:$PATH

# 安装 Python 包
pip install $TRT_PATH/python/tensorrt-8.6.1-cp39-none-linux_x86_64.whl

# 验证安装
python -c "import tensorrt as trt; print(trt.__version__); print(trt.utils.get_builder_version())"

验证输出应该是 8.6.1 之类的版本号。如果报错说找不到库,八成是 LD_LIBRARY_PATH 没设对。

我曾经踩过的坑: 用 pip 安装了 TensorRT,结果运行时一直报算子不支持。查了半天发现 pip 版本是 8.4,而我需要的 8.6 版本的新特性。从那以后,我都是手动下载 tar 包安装,版本控制更精准。

2.3.2 NCNN 安装

NCNN 是腾讯开源的移动端推理引擎,对 ARM 平台优化得很好。安装方式是从源码编译:

# 克隆源码
git clone https://github.com/Tencent/ncnn.git
cd ncnn
git submodule update --init

# 编译
mkdir build && cd build
cmake -DCMAKE_BUILD_TYPE=Release -DNCNN_VULKAN=ON ..
make -j$(nproc)

# 安装
make install

# 验证(运行自带测试)
./tools/ncnn2table --help

如果编译成功,ncnn2table 工具就能正常使用了。Vulkan 选项建议打开,这样可以利用 GPU 加速。

2.3.3 OpenVINO 安装

OpenVINO 是 Intel 推出的推理引擎,对 Intel CPU、集成显卡、VPU 都有优化。安装方式推荐用 pip:

# 安装 OpenVINO 2023.3
pip install openvino==2023.3.0

# 验证安装
python -c "import openvino as ov; print(ov.__version__); core = ov.Core(); print(core.available_devices)"

验证输出应该显示版本号和可用设备列表,比如 ['CPU', 'GPU']。如果只有 CPU,说明 GPU 驱动没装好,但不影响基本使用。

2.4 工具链验证与整体架构

所有工具装完后,我建议你跑一个完整的端到端验证:从 PyTorch 模型 → ONNX → TensorRT/NCNN/OpenVINO,确保整个链路是通的。

下面这张图展示了整个工具链的架构和转换流程:

模型压缩与部署工具链架构 训练框架 PyTorch / TensorFlow 中间表示层 ONNX(通用模型格式) 推理引擎层 TensorRT / NCNN / OpenVINO 硬件平台层 NVIDIA GPU / ARM CPU / Intel CPU 各层作用说明 1. 训练框架:模型训练与导出 支持导出 ONNX 格式 2. 中间表示层:模型格式统一 解决框架间算子差异 3. 推理引擎:模型优化与加速 TensorRT:NVIDIA GPU 优化 NCNN:ARM 移动端优化 OpenVINO:Intel 平台优化 4. 硬件平台:实际部署环境 根据硬件选择对应引擎

从这张图可以看得很清楚:训练框架在最上层,中间通过 ONNX 统一格式,下层对接不同的推理引擎,最终适配不同硬件。每一层都有自己独立的安装和验证流程,但整体是串在一起的。

2.5 常见问题与避坑指南

环境搭建过程中,有几个问题几乎每个人都会遇到。我把它们列出来,你遇到了直接对照排查:

问题现象 可能原因 解决方案
PyTorch 检测不到 GPU CUDA 版本不匹配或驱动未安装 运行 nvidia-smi 检查驱动,重新安装对应 CUDA 版本的 PyTorch
ONNX 导出报算子不支持 模型使用了 ONNX 不支持的算子 用 torch.onnx.export 的 operator_export_type 参数降级,或替换算子
TensorRT 转换失败 ONNX 模型中有动态 shape 或不受支持的层 先用 onnx-simplifier 简化模型,再手动处理不支持层
NCNN 编译报错 缺少依赖库或 CMake 版本过低 安装 protobuf、opencv 等依赖,升级 CMake 到 3.13+
OpenVINO 找不到设备 GPU 驱动未安装或权限不足 安装 Intel GPU 驱动,或添加用户到 render 组
我的建议: 环境搭建完成后,把每个工具的验证命令写成一个脚本,每次换机器或重装系统时直接跑一遍。我自己的 env_check.sh 脚本已经迭代了十几个版本,现在换新机器基本半小时就能搞定全套环境。

好了,环境搭建这部分就到这里。你跟着把每个工具装好、验证通过,后面做模型转换和推理优化的时候就不会被环境问题卡住了。记住,工具链的版本一致性是重中之重,差一个小版本都可能让你多花半天排查时间。