第2章：环境准备——把工具箱先搭稳了

说实话，做模型剪枝最怕什么？不是剪完精度掉得厉害，而是环境没配好，折腾半天连个demo都跑不起来。我见过太多同学卡在CUDA版本冲突上，一卡就是一整天。所以这一章，咱们先把地基打牢。

我会带你装好三样东西：PyTorch（深度学习框架）、Ultralytics（YOLOv11的官方库）、torch-pruning（剪枝工具）。另外CUDA环境和数据集路径也得配好。别急，一步步来。

2.1 环境总览：我们要装什么？

先看一张整体流程图，让你心里有个谱：

2.2 安装PyTorch——选对版本省一半力气

PyTorch是整个剪枝流程的基石。我个人习惯用pip安装，干净利落。但有个坑——CUDA版本必须和PyTorch匹配。

先查一下你的显卡驱动支持的最高CUDA版本：

nvidia-smi
# 看右上角 "CUDA Version: 12.1" 之类的

然后去PyTorch官网（pytorch.org）选对应的安装命令。举个例子，CUDA 12.1就用：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

装完验证一下：

import torch
print(torch.__version__)        # 比如 2.1.0+cu121
print(torch.cuda.is_available()) # 必须 True

2.3 安装Ultralytics——YOLOv11的官方入口

Ultralytics库把YOLOv11的训练、推理、导出都封装好了。安装很简单：

pip install ultralytics

但注意，版本别太老。我建议用8.0.200以上：

pip install ultralytics==8.0.200

验证一下：

from ultralytics import YOLO
model = YOLO('yolo11n.pt')  # 会自动下载预训练权重
print(model)                 # 能看到模型结构

2.4 安装torch-pruning——剪枝的核心武器

torch-pruning是我个人非常喜欢的一个剪枝库，它支持结构化剪枝和非结构化剪枝。安装方式：

pip install torch-pruning

或者从源码安装（推荐，功能最新）：

git clone https://github.com/VainF/Torch-Pruning.git
cd Torch-Pruning
pip install -e .

验证一下：

import torch_pruning as tp
print(tp.__version__)  # 比如 1.3.0

你想想看，剪枝的本质就是在模型里插入“掩码”，让某些权重不参与计算。torch-pruning帮我们自动处理了这些细节。

2.5 配置CUDA环境——别让GPU闲着

CUDA环境说白了就是让PyTorch能调用GPU。除了驱动匹配，还有几个要点：

• cuDNN：深度学习加速库，一般PyTorch自带，不用单独装
• 环境变量：有时候需要手动指定CUDA路径

export CUDA_HOME=/usr/local/cuda-12.1
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

验证GPU是否被PyTorch识别：

import torch
print(torch.cuda.device_count())  # GPU数量
print(torch.cuda.get_device_name(0))  # 比如 "NVIDIA GeForce RTX 3090"

2.6 配置数据集路径——让YOLO找到数据

YOLOv11用YAML文件管理数据集路径。以COCO为例，创建一个 coco.yaml：

# coco.yaml
path: /path/to/coco  # 数据集根目录
train: images/train2017
val: images/val2017
test: images/test2017

nc: 80  # 类别数
names: ['person', 'bicycle', 'car', ...]  # 类别名

如果是自定义数据集，格式一样。我个人习惯把数据集放在 ~/datasets/ 下，方便管理。

验证数据集路径是否正确：

from ultralytics import YOLO
model = YOLO('yolo11n.pt')
results = model.val(data='coco.yaml')  # 能跑通就说明路径对了

2.7 环境验证——跑通第一个demo

所有东西装好后，跑一个完整的训练-剪枝-验证流程，确保环境没问题：

from ultralytics import YOLO
import torch_pruning as tp

# 1. 加载模型
model = YOLO('yolo11n.pt')

# 2. 简单剪枝测试（非结构化，剪掉20%）
model.model = tp.prune_conv_non_structured(model.model, amount=0.2)

# 3. 验证
results = model.val(data='coco.yaml')
print(f"mAP@0.5: {results.box.map50:.3f}")

如果这一步没报错，恭喜你，环境配好了！

2.8 常见问题速查表

问题	原因	解决方法
torch.cuda.is_available() 为 False	CUDA版本不匹配或没装GPU版PyTorch	重装对应CUDA版本的PyTorch
导入ultralytics报错	版本太老或依赖冲突	pip install --upgrade ultralytics
torch-pruning导入失败	没装或版本不兼容	从源码安装：pip install -e .
数据集路径报错	YAML中路径写错了	用路径，检查文件夹是否存在

嗯，环境准备就到这儿。别嫌麻烦，这一步稳了，后面剪枝、微调、精度恢复才能顺畅。我个人每次开新项目都会花半小时把环境彻底验证一遍，省得后面出幺蛾子。