2. 环境搭建：PyTorch与CUDA环境配置、YOLOv11源码下载与依赖安装、验证训练流程是否跑通

说实话，很多同学做模型剪枝，最后模型没剪好，反而是环境先崩了。

我见过太多人花了一周配环境，结果训练都跑不起来。所以这一章，咱们先把地基打牢。你跟着我的步骤走，半小时内让YOLOv11跑起来。

2.1 硬件与软件版本选型

先看看你的硬件。我个人的建议是：

嗯，这里要注意。显存6GB跑YOLOv11s勉强可以，但如果你要做剪枝后的微调，建议至少8GB。我在项目里吃过亏，用4GB的卡跑剪枝模型，batch size设到2都爆显存。

2.2 CUDA与cuDNN安装

先确认你的驱动版本。打开终端：

nvidia-smi

看右上角的CUDA Version。比如显示12.1，那你就装CUDA 12.1对应的PyTorch版本。

我个人习惯用conda管理环境，干净利落：

conda create -n yolov11 python=3.10 -y
conda activate yolov11

然后装PyTorch。去PyTorch官网选对应版本。我常用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

python -c "import torch; print(torch.cuda.is_available())"

2.3 YOLOv11源码下载

YOLOv11官方仓库在GitHub上。我建议你fork一份到自己的账号，方便后面改代码做剪枝。

git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics

然后安装依赖：

pip install -r requirements.txt

这里有个小技巧。我一般会额外装几个后面剪枝会用到的库：

pip install thop  # 计算FLOPs
pip install onnx onnxruntime  # 模型导出
pip install nvidia-ml-py3  # 监控显存

2.4 验证训练流程

环境配好了，得跑个demo验证一下。别一上来就训大数据集，先用COCO128这个小数据集跑通流程。

下载COCO128：

# 脚本会自动下载到 datasets/ 目录
yolo train data=coco128.yaml model=yolov11n.pt epochs=1

如果报错，别慌。我整理了几个常见问题：

yolo train data=coco128.yaml model=yolov11n.pt epochs=1 verbose=True

2.5 验证剪枝环境

既然咱们要做剪枝，还得确认几个关键库能不能用。写个脚本跑一下：

import torch
import torch.nn.utils.prune as prune
from ultralytics import YOLO

# 加载模型
model = YOLO('yolov11n.pt')
print("模型加载成功")

# 检查是否能访问模型参数
for name, module in model.model.named_modules():
    if hasattr(module, 'weight'):
        print(f"可剪枝模块: {name}, 权重形状: {module.weight.shape}")
        break

print("剪枝环境验证通过 ✅")

如果输出正常，说明你的环境已经准备好做剪枝了。

2.6 知识结构总览

这一章的内容，说白了就是三步走。我画了个图帮你理清思路：

2.7 验证结果确认清单

最后，我给你一个检查清单。每完成一项打个勾：

• ✅ nvidia-smi 显示GPU信息
• ✅ torch.cuda.is_available() 返回 True
• ✅ YOLOv11源码下载成功，pip list 能看到 ultralytics
• ✅ yolo train epochs=1 能跑完一个epoch
• ✅ 剪枝环境验证脚本无报错