第二章 环境搭建与工具链：PyTorch、Transformers、ONNX、TensorRT 配置

做模型蒸馏和压缩，说白了就是跟四个工具打交道：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

2.3 Transformers 库安装

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

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

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

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

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

2.6 版本兼容性速查表

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

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

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 环境，互不干扰。

环境搭好了，后面的事情就顺了。接下来我们会进入真正的蒸馏流程——从数据准备到模型训练，每一步我都会带着你走一遍。