📚 课程目录第2章环境搭建与工具链

2. 环境搭建与工具链

做模型蒸馏和压缩，说白了就是跟四个工具打交道：PyTorch 负责训练，Transformers 提供预训练骨架，ONNX 做中间格式转换，TensorRT 做最终推理加速。这四个东西的版本搭配和环境配置，我见过太多人在这里卡好几天。

这一章咱们就把环境彻底捋顺。我会把每个工具的版本选择、安装命令、常见坑点都过一遍。你跟着做一遍，后面二十多章就不会再被环境问题打断。

2.1 整体工具链架构

先看一张整体流程图，搞清楚每个工具在蒸馏压缩流程里扮演什么角色：

这张图我建议你保存下来。每次做到对应步骤时，回来看一眼，就知道当前在哪个环节。

2.2 PyTorch 环境配置

PyTorch 是整个流程的基础。我个人习惯用 Conda 管理环境，干净、可复现。

2.2.1 创建虚拟环境

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

Python 版本我推荐 3.9 或 3.10。3.11 虽然新，但有些 ONNX 算子兼容性还没完全跟上。

2.2.2 安装 PyTorch

去 PyTorch 官网选对应 CUDA 版本的命令。以 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

我的建议： 别用 conda install pytorch，那个版本经常滞后。用 pip 装更可控。装完后跑一下 python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())" 确认 GPU 可用。

2.3 Transformers 库安装

Transformers 库提供预训练模型和蒸馏相关的工具类。我建议装最新稳定版：

pip install transformers==4.36.0 datasets==2.16.0 accelerate==0.25.0

这里我踩过一个坑：accelerate 版本太低会导致多卡蒸馏时数据并行报错。所以这三个库最好一起装，版本对齐。

库名	推荐版本	作用
transformers	≥ 4.36.0	模型加载、蒸馏 Trainer
datasets	≥ 2.16.0	数据加载与预处理
accelerate	≥ 0.25.0	分布式训练加速

2.4 ONNX 与 ONNX Runtime 配置

ONNX 是模型从 PyTorch 到 TensorRT 的桥梁。这一步配置不好，后面导出会各种报错。

2.4.1 安装 ONNX 工具包

pip install onnx==1.15.0 onnxruntime-gpu==1.16.3 onnxconverter-common==1.13.0

注意： 千万别装 onnxruntime（CPU 版）！一定要装 onnxruntime-gpu。我曾经因为装错版本，推理速度反而比 PyTorch 还慢，排查了半天才发现是 CPU 版在跑。

2.4.2 验证 ONNX 导出

装完后跑一个简单测试：

import torch
import torch.onnx

dummy_input = torch.randn(1, 3, 224, 224).cuda()
model = torchvision.models.resnet18().cuda()
torch.onnx.export(model, dummy_input, "test.onnx", opset_version=17)
print("ONNX export success")

如果这一步报错，大概率是 opset 版本不匹配。我一般用 opset 17，兼容性最好。

2.5 TensorRT 环境搭建

TensorRT 是压轴环节。它的安装稍微麻烦一点，因为需要跟 CUDA 版本严格对应。

2.5.1 确认 CUDA 版本

nvcc --version

记下输出的 CUDA 版本，比如 11.8 或 12.1。然后去 NVIDIA 官网下载对应的 TensorRT 版本。

2.5.2 安装 TensorRT

我推荐用 tar 包安装，比 deb 包更灵活：

# 以 TensorRT 8.6.1 CUDA 11.8 为例
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/8.6.1/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz

# 添加环境变量
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/TensorRT-8.6.1.6/lib
export PATH=$PATH:/path/to/TensorRT-8.6.1.6/bin

# 安装 Python 包
pip install /path/to/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp39-none-linux_x86_64.whl

  关键点： TensorRT 的 Python 包名里包含 Python 版本号（cp39 表示 Python 3.9）。选错版本 import 会直接报错。装完后跑 python -c "import tensorrt as trt; print(trt.__version__)" 验证。

2.6 版本兼容性速查表

我把常用版本组合整理成了一张表，你直接照着配就行：

组件	组合一（推荐）	组合二
CUDA	11.8	12.1
PyTorch	2.1.0	2.2.0
Transformers	4.36.0	4.38.0
ONNX	1.15.0	1.16.0
ONNX Runtime GPU	1.16.3	1.17.0
TensorRT	8.6.1	9.0.1

我个人更推荐组合一，因为踩过的坑少，社区资料也最全。

2.7 常见问题与避坑指南

CUDA 版本不匹配： 装 PyTorch 时 CUDA 版本跟系统不一致，导致 torch.cuda.is_available() 返回 False。解决办法：用 nvidia-smi 看驱动支持的 CUDA 版本，然后选对应版本的 PyTorch。
ONNX 导出时算子不支持： 有些 PyTorch 操作（如 torch.einsum）ONNX 不支持。我曾经在导出 BERT 模型时遇到这个问题，后来改用 torch.bmm 替代才解决。
TensorRT 引擎构建失败： 多半是动态 shape 设置不对。建议先用固定 batch size 和序列长度导出 ONNX，等引擎构建成功后再改动态。
环境冲突： 不同项目用不同版本的 PyTorch 和 TensorRT，容易打架。我的习惯是每个项目一个独立的 Conda 环境，互不干扰。

一个小技巧： 装完所有工具后，写一个 requirements.txt 把版本固定下来。下次重建环境时直接 pip install -r requirements.txt，省时省力。