2、环境搭建与工具准备：PyTorch、TensorFlow、ONNX、TensorRT等框架的安装配置，以及实验数据集和预训练模型的选择
模型压缩（剪枝、量化、蒸馏）的工程落地，首先依赖一套稳定、可复现且性能达标的软件环境。本章聚焦于主流深度学习框架（PyTorch、TensorFlow）、中间表示（ONNX）以及推理优化引擎（TensorRT）的安装配置要点，同时给出实验数据集与预训练模型的选型策略。环境搭建并非简单的“pip install”，而是涉及CUDA版本、cuDNN、Python解释器、框架编译选项以及硬件算力之间的复杂耦合。一个错误的版本组合可能导致训练速度骤降、算子不支持甚至显存溢出等诡异问题。因此，我们需要从系统层面理解各组件的关系，再动手实践。
2.1 核心原理：理解框架、中间表示与推理引擎的层级关系
在搭建环境之前，必须明确各工具在模型压缩与部署流程中的定位。PyTorch和TensorFlow是训练与剪枝/量化/蒸馏实验的主要载体，它们提供了自动求导、优化器和丰富的算子库。ONNX（Open Neural Network Exchange）作为中间表示层，负责将训练好的模型从特定框架中“解放”出来，转换为框架无关的计算图描述。TensorRT则是NVIDIA GPU上的高性能推理引擎，它对ONNX或TensorFlow模型进行层融合、精度校准（如INT8量化）和内核自动调优，从而显著降低延迟。
这种分层架构决定了环境配置的优先级：首先确保GPU驱动和CUDA运行时正确，然后安装深度学习框架，接着配置ONNX导出与转换工具链，最后根据部署目标安装TensorRT。任何一层的版本不匹配，都会导致后续步骤失败。例如，PyTorch 2.x默认使用CUDA 11.8或12.1编译，如果系统只安装了CUDA 11.4，则无法启用GPU加速。
2.2 判断方法与操作流程：从零搭建一套可用的压缩实验环境
以下操作流程以Ubuntu 20.04/22.04 + NVIDIA GPU（算力≥7.5，如RTX 30/40系列）为例。请根据实际硬件调整版本号。
步骤一：安装GPU驱动与CUDA
# 检查GPU型号
nvidia-smi
# 推荐安装驱动版本 ≥ 525，然后安装CUDA 11.8或12.1
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run --toolkit --silent --override
# 配置环境变量
echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
nvcc -V # 验证CUDA版本步骤二：创建虚拟环境并安装PyTorch
conda create -n compress python=3.9 -y
conda activate compress
# 安装PyTorch（根据CUDA版本选择）
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
# 验证GPU可用
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"步骤三：安装ONNX工具链
pip install onnx==1.14.0 onnxruntime-gpu==1.16.0 onnxsim
# 验证ONNX导出
python -c "import onnx; print(onnx.__version__)"步骤四：安装TensorRT（可选，用于部署验证）
# 下载TensorRT 8.6 GA，使用tar包解压
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/TensorRT-8.6.1.6/lib' >> ~/.bashrc
pip install /path/to/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp39-none-linux_x86_64.whl
# 验证
python -c "import tensorrt as trt; print(trt.__version__)"| 组件 | 推荐版本 | 关键注意事项 |
|---|---|---|
| NVIDIA Driver | ≥ 525.60 | 驱动版本必须高于CUDA Toolkit要求的最低版本 |
| CUDA Toolkit | 11.8 / 12.1 | PyTorch 2.x对CUDA 11.8支持最稳定 |
| PyTorch | 2.1.0 | 支持torch.compile，对量化感知训练有优化 |
| ONNX | 1.14+ | 配合onnxsim简化模型结构 |
| TensorRT | 8.6 GA | 需与CUDA版本严格对应 |
2.3 案例拆解：剪枝实验中的环境问题排查与数据集选择
背景与约束：某团队需要对ResNet-50进行结构化剪枝，目标是将模型压缩至原始体积的30%，并在TensorRT上实现INT8量化推理。团队使用PyTorch 1.12 + CUDA 11.3环境，但在导出ONNX时发现部分算子（如F.relu6）不支持。
初始现象：onnx.export()报错“Unsupported operator: relu6”。进一步检查发现，PyTorch 1.12的relu6导出需要额外注册算子。
分析判断：问题根源在于PyTorch与ONNX的算子版本映射不完整。解决方案有两种：一是升级PyTorch至2.0+，其ONNX导出器已支持更多算子；二是使用onnx-simplifier进行图优化，将relu6转换为relu + 常数缩放。
实际操作步骤：
- 升级PyTorch至2.1.0，并重新导出模型。
- 使用
onnxsim对导出的ONNX模型进行简化，消除冗余节点。 - 在TensorRT中构建引擎时，设置
FP16或INT8精度，并利用trtexec工具验证延迟。
结果对比：升级后，ONNX导出成功，TensorRT INT8推理延迟从原始PyTorch的12ms降至2.3ms，精度损失仅0.8%。
风险与代价：升级PyTorch可能导致旧代码API不兼容，需回归测试训练流程。此外，INT8校准需要额外的校准数据集（如ImageNet的1000张子集），增加了数据准备成本。
复盘经验：环境搭建时，应优先选择“训练框架最新稳定版 + ONNX最新版 + TensorRT匹配版本”的组合，避免使用过旧版本。同时，数据集选择应遵循“与预训练模型同分布”的原则，例如使用ImageNet预训练权重时，校准集也应从ImageNet中采样。
2.4 实验数据集与预训练模型的选择策略
数据集和预训练模型是压缩实验的“输入质量”决定因素。对于分类任务，CIFAR-10/100适合快速验证剪枝算法；ImageNet（ILSVRC2012）是标准基准，但训练成本高。对于检测任务，COCO 2017是主流选择。预训练模型方面，PyTorch官方torchvision库提供了ResNet、MobileNet、EfficientNet等权重，TensorFlow Hub也有对应模型。
预训练模型下载与校验
import torchvision.models as models
# 下载ResNet50预训练权重
resnet50 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
# 保存为本地文件
torch.save(resnet50.state_dict(), 'resnet50_imagenet.pth')
# 校验模型输出
import torch
dummy = torch.randn(1, 3, 224, 224)
out = resnet50(dummy)
assert out.shape == (1, 1000)2.5 常见误区与总结复盘
- 误区一：盲目追求最新CUDA版本。实际上，PyTorch和TensorRT对CUDA版本有严格的上限要求，过新的CUDA可能导致框架无法识别GPU。
- 误区二：忽略cuDNN的安装。cuDNN是卷积运算的加速库，缺失时训练速度会下降数倍。
- 误区三：在同一个环境中混用TensorFlow和PyTorch。两者的依赖冲突难以调和，建议使用独立虚拟环境。
- 误区四：使用与预训练模型输入尺寸不一致的数据集。例如，将ImageNet预训练模型直接用于32x32的CIFAR-10，需要调整输入分辨率或修改网络结构。