2、环境搭建与工具链:Python环境配置、PyTorch安装、ONNX与TensorRT简介、Ultralytics库安装
做模型压缩和推理加速,第一步不是调参,而是把环境搭稳。我见过太多同学花了两天调代码,最后发现是CUDA版本没对齐——这种坑,踩一次就够了。
这一章,咱们把Python、PyTorch、ONNX、TensorRT和Ultralytics一条龙装好。嗯,我会把每个环节的坑都指出来。
2.1 Python环境配置
我个人习惯用Miniconda而不是Anaconda。为什么?轻量。Anaconda预装了一堆你用不到的包,Miniconda只带核心,干净利落。
推荐版本:Python 3.10。PyTorch 2.x对3.10支持最好,3.11以上有些旧版ONNX opset会报错。
安装步骤很简单:
- 去Miniconda官网下载对应系统的安装包
- 安装时勾选"Add to PATH"(Windows用户注意)
- 打开终端,验证:
conda --version
然后创建一个专属环境:
conda create -n yolov11 python=3.10 -y
conda activate yolov11
为什么要单独建环境?我在项目中吃过亏——不同项目依赖打架,最后只能重装系统。隔离环境,这是第一道保险。
2.2 PyTorch安装
PyTorch是YOLOv11的基石。选版本有个原则:稳定版 + 对应CUDA。
| 场景 | 推荐命令 | 说明 |
|---|---|---|
| 有NVIDIA GPU(推荐) | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 |
CUDA 12.1,兼容RTX 30/40系列 |
| 无GPU / Mac | pip install torch torchvision torchaudio |
CPU版本,MPS加速需额外配置 |
| 老显卡(GTX 10系列) | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 |
CUDA 11.8,兼容性更好 |
避坑:我曾经在服务器上装了PyTorch 2.0 + CUDA 11.7,结果TensorRT死活不认。后来发现TensorRT 8.6只支持CUDA 11.8+。所以,先查TensorRT版本,再定CUDA版本。
装完后验证一下:
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"
如果输出True,恭喜,GPU环境通了。
2.3 ONNX与TensorRT简介
这两个工具是模型压缩和加速的核心。说白了:
- ONNX:模型的中转格式。不管你是PyTorch、TensorFlow还是PaddlePaddle,都能转成ONNX。它像个通用翻译器。
- TensorRT:NVIDIA的推理优化器。它把ONNX模型吃进去,吐出一个更小、更快的引擎。
安装ONNX很简单:
pip install onnx onnxruntime-gpu
TensorRT稍微麻烦点。我建议用tar包安装,而不是pip——因为pip版本经常滞后。
- 去NVIDIA官网下载TensorRT 8.6+(对应你的CUDA版本)
- 解压后,把lib路径加到
LD_LIBRARY_PATH - 安装Python wheel:
pip install tensorrt-*-cp310-none-linux_x86_64.whl
核心概念:ONNX负责"格式统一",TensorRT负责"极致优化"。两者配合,模型推理速度能提升2~5倍。
验证TensorRT安装:
python -c "import tensorrt as trt; print(trt.__version__)"
2.4 Ultralytics库安装
Ultralytics是YOLOv11的官方库。它封装了训练、验证、导出、推理全流程。
安装命令:
pip install ultralytics
它会自动拉依赖,包括opencv-python、matplotlib、pandas等。如果你网络慢,可以加个国内镜像:
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple
装完后跑个测试:
yolo predict model=yolo11n.pt source='https://ultralytics.com/images/bus.jpg'
如果看到检测结果,说明环境全通了。
我的习惯:每次装完Ultralytics,我会先跑一次
yolo export model=yolo11n.pt format=onnx,验证ONNX导出链路是否正常。这一步能提前发现opset不兼容的问题。
2.5 整体工具链关系图
下面这张图,把整个工具链的流转关系画清楚了。你想想看,从训练到部署,每一步都在干什么:
2.6 环境验证清单
装完所有东西后,我建议你跑一遍这个清单。别嫌麻烦——我在项目里吃过"装完就跑,跑完就崩"的亏。
| 检查项 | 命令 | 预期结果 |
|---|---|---|
| Python版本 | python --version |
3.10.x |
| PyTorch版本 | python -c "import torch; print(torch.__version__)" |
2.x |
| CUDA可用 | python -c "import torch; print(torch.cuda.is_available())" |
True |
| ONNX版本 | python -c "import onnx; print(onnx.__version__)" |
1.14+ |
| TensorRT版本 | python -c "import tensorrt as trt; print(trt.__version__)" |
8.6+ |
| Ultralytics | yolo version |
8.3+ |
曾经踩过的坑:有一次我装好了所有东西,但
torch.cuda.is_available()一直返回False。查了半天,发现是CUDA Toolkit版本和显卡驱动不匹配。记住:驱动版本 ≥ CUDA Toolkit版本。用nvidia-smi看驱动支持的CUDA版本,再装对应的PyTorch。
环境搭好了,后面的事就顺了。接下来我们会用这套工具链,一步步把YOLOv11模型压缩到极致。