📚 课程目录第2章环境搭建

2. 环境搭建

说实话，环境搭建这事儿看着简单，但我在项目里见过太多人在这上面栽跟头。有的同学装了半天，最后发现CUDA版本对不上；有的明明配好了，一跑就报错。嗯，咱们这一章就把这事儿捋清楚。

2.1 整体环境概览

先看看咱们需要哪些东西。我习惯把环境分成三层：

层级	组件	作用
基础层	Python 3.8+、CUDA 11.8+、cuDNN	提供运行底座和GPU加速能力
框架层	PyTorch 2.0+、ONNX 1.14+	模型训练、导出、中间表示
优化层	TensorRT 8.6+、ONNX Runtime、NVIDIA AMMO	剪枝、量化、推理加速

说白了，这三层缺一不可。底层不对，上层全白搭。我建议你先把基础层搞定，再往上走。

核心原则：所有组件的版本必须互相兼容。别图新，也别图省事。我见过有人装了CUDA 12.0，结果PyTorch 2.0不支持，折腾了两天。

2.2 基础环境安装

2.2.1 Python环境

我个人习惯用conda管理环境。干净、隔离、好回滚。你想想看，要是直接在系统Python里装一堆包，哪天搞崩了，哭都来不及。

# 创建独立环境
conda create -n yolov11_opt python=3.10 -y
conda activate yolov11_opt

# 升级pip
pip install --upgrade pip setuptools wheel

小技巧：环境名我习惯用项目名_用途的格式，比如yolov11_opt。这样项目多了也不会乱。

2.2.2 CUDA与cuDNN

这里有个坑。CUDA版本不是越新越好，得看PyTorch和TensorRT支持什么。我目前在用的组合是CUDA 11.8 + cuDNN 8.9，稳得很。

# 查看CUDA版本
nvcc --version

# 查看cuDNN版本
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

注意：如果你用conda装PyTorch，它会自动帮你装cudatoolkit。但TensorRT需要系统级的CUDA。我曾经图省事只装了conda版，结果TensorRT死活找不到CUDA库。嗯，后来老老实实装了系统版。

2.3 框架层安装

2.3.1 PyTorch

PyTorch 2.0以上版本对torch.compile支持更好，剪枝量化时能省不少事。我建议直接上2.1或2.2。

# CUDA 11.8版本
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118

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

如果输出True，说明GPU可用。如果输出False，别慌，八成是CUDA版本没对上。

2.3.2 ONNX与ONNX Runtime

ONNX是模型转换的中间桥梁。剪枝后的模型要转成ONNX，再喂给TensorRT。ONNX Runtime则用来做推理验证。

pip install onnx==1.14.0 onnxruntime-gpu==1.16.0

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

经验之谈：ONNX版本别追最新。我遇到过1.15版本导出时某些算子不兼容，回退到1.14就没事了。稳定比新功能重要。

2.4 优化层安装

2.4.1 TensorRT

TensorRT是咱们做量化的核心工具。安装方式有两种：pip安装和tar包安装。我推荐tar包方式，能自定义安装路径，方便管理。

# 下载TensorRT 8.6.1 tar包
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-cp310-none-linux_x86_64.whl

避坑指南：我曾经直接把TensorRT的lib路径加到/etc/ld.so.conf里，结果跟系统其他库冲突了。后来改用LD_LIBRARY_PATH环境变量，只在当前终端生效，安全多了。

2.4.2 NVIDIA AMMO（模型优化工具）

AMMO是NVIDIA出的模型优化工具包，专门做剪枝和量化。它基于PyTorch，跟咱们的YOLOv11配合得很好。

pip install nvidia-ammo

# 验证
python -c "import ammo.torch.quantization as atq; print('AMMO ready')"

2.5 完整环境验证

装完所有东西后，我习惯跑一个完整的验证脚本。确保所有组件能协同工作。

import torch
import onnx
import tensorrt as trt

print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"ONNX: {onnx.__version__}")
print(f"TensorRT: {trt.__version__}")

# 简单测试：创建一个小模型并转ONNX
model = torch.nn.Linear(10, 5).cuda()
dummy_input = torch.randn(1, 10).cuda()
torch.onnx.export(model, dummy_input, "test.onnx")
print("ONNX export: OK")

# 测试TensorRT
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
print("TensorRT builder: OK")

如果所有输出都正常，恭喜你，环境搭好了。

2.6 环境结构总览

为了让你更直观地理解各组件的关系，我画了张图：

从下往上看，基础层支撑框架层，框架层再支撑优化层。每一层都依赖下一层提供的能力。你想想看，要是CUDA没装好，PyTorch就调不了GPU；PyTorch没装好，AMMO就没法做量化。所以，按顺序来，别跳步。

2.7 常见问题与解决方案

问题	原因	解决方案
`torch.cuda.is_available()`返回False	CUDA版本不匹配或驱动未安装	检查`nvidia-smi`输出，确认驱动版本；重新安装匹配的PyTorch
ONNX导出时报算子不支持	ONNX版本过新或过旧	回退到ONNX 1.14，或升级PyTorch到2.1+
TensorRT构建engine失败	CUDA库路径未正确设置	检查`LD_LIBRARY_PATH`是否包含TensorRT的lib目录
AMMO导入报错	依赖的PyTorch版本不兼容	使用PyTorch 2.0-2.2，不要用2.3+

我的习惯：每次装完一个组件，都跑一遍验证脚本。别等全装完了再查问题。那样你根本不知道是哪个环节出的错。

好了，环境搭完了。下一步咱们就可以开始真正的剪枝和量化工作了。记住，环境是基础，基础不牢，地动山摇。花点时间把这一步做扎实，后面会顺很多。