2. 环境准备

做模型剪枝，说白了就是给 YOLO 做“瘦身手术”。手术刀再锋利，手术台没搭好也是白搭。这一章，我就带你把这个台子搭稳了。

环境配置这事，看着琐碎，但坑不少。我见过太多同学卡在第一步，CUDA 版本对不上、Python 包冲突，折腾半天还没开始跑训练。嗯，咱们别走弯路，一次搞定。

2.1 Python 虚拟环境：隔离的艺术

我个人习惯，每个项目都建独立的虚拟环境。为什么？你想想看，YOLOv11 依赖的 PyTorch 版本、torchvision 版本，跟其他项目可能打架。隔离了，省心。

我用的是 conda，当然 venv 也行。这里我推荐 conda，因为后面装 CUDA 相关包时，它处理依赖更干净。

# 创建环境，指定 Python 3.10（YOLOv11 官方推荐）
conda create -n yolov11_prune python=3.10 -y

# 激活环境
conda activate yolov11_prune

# 验证 Python 版本
python --version
# 输出应为：Python 3.10.x

2.2 PyTorch 与 CUDA：算力的心脏

这一步最容错。我曾经在项目里因为 CUDA 版本装错，白花了半天排查。记住一个原则：PyTorch 版本必须匹配 CUDA 驱动版本。

先看看你的显卡支持什么 CUDA 版本：

# 查看 CUDA 驱动版本（不是 toolkit 版本）
nvidia-smi

# 输出示例：
# +-----------------------------------------------------------------------------+
# | NVIDIA-SMI 545.23.08    Driver Version: 545.23.08    CUDA Version: 12.3     |
# +-----------------------------------------------------------------------------+

看到 CUDA Version: 12.3，说明驱动支持最高 CUDA 12.3。然后去 PyTorch 官网选对应版本。

我推荐用 pip 安装，干净利落：

# 以 CUDA 12.1 为例（稳定且兼容性好）
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121

# 验证安装
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"
# 输出应为：
# 2.1.0
# True

2.3 YOLOv11 源码下载与依赖安装

环境搭好了，接下来拉代码。我建议直接从官方仓库克隆，别用第三方打包的，容易出问题。

# 克隆 YOLOv11 源码
git clone https://github.com/ultralytics/ultralytics.git

# 进入目录
cd ultralytics

# 安装依赖（核心）
pip install -r requirements.txt

# 额外安装剪枝可能用到的工具
pip install thop  # 计算 FLOPs
pip install onnx  # 模型导出

依赖装完后，我习惯看一眼 requirements.txt，确认关键包的版本：

包名	最低版本	说明
torch	≥1.8.0	我们装的是 2.1.0，够用
torchvision	≥0.9.0	与 torch 版本匹配
numpy	≥1.18.0	基础运算
opencv-python	≥4.1.0	图像处理
matplotlib	≥3.2.2	可视化

2.4 验证环境：跑通第一个 Demo

环境装完，必须跑个 Demo 验证。我习惯用官方预训练模型做一次推理，确保整个链路通顺。

# 下载预训练模型（会自动下载）
yolo predict model=yolov11n.pt source='https://ultralytics.com/images/bus.jpg'

# 或者用 Python 脚本
from ultralytics import YOLO

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

# 推理
results = model('https://ultralytics.com/images/bus.jpg')

# 打印结果
for r in results:
    print(r.boxes)  # 检测框信息

如果看到类似下面的输出，说明环境没问题：

image 1/1 /path/to/bus.jpg: 640x480 4 persons, 1 bus, 1 stop sign, 87.2ms
Speed: 3.2ms preprocess, 87.2ms inference, 1.1ms postprocess per image at shape (1, 3, 640, 480)

2.5 本章知识体系总览

下面这张图，把整个环境准备流程串起来了。你可以把它贴在工位上，每次配环境时对照着来。

2.6 常见问题速查表

最后，我把环境配置中最容易踩的坑整理成一张表。遇到问题先查这里，能省不少时间。

问题现象	可能原因	解决方案
torch.cuda.is_available() 返回 False	CUDA 驱动版本过低或 PyTorch 版本不匹配	运行 nvidia-smi 查看驱动版本，安装对应 PyTorch
ImportError: No module named 'ultralytics'	未安装或安装路径不对	在 ultralytics 目录下运行 pip install -e .
推理时显存不足	模型太大或 batch size 太大	换小模型（如 yolov11n）或降低 batch size
依赖版本冲突	多个项目共用同一环境	使用虚拟环境隔离，或 pip install 包名==指定版本