3. 环境准备：搭建剪枝与调优所需的Python环境

好，咱们正式开始动手。这一章我带你把整个实验环境搭起来。别小看这一步，我在项目里见过太多人因为环境冲突浪费一整天。说白了，环境配好了，后面才能跑得顺。

3.1 环境总览：我们需要什么？

先看看我们要装哪些东西。我习惯把依赖分成三层：核心框架、剪枝工具、辅助库。这样出了问题也好排查。

3.2 创建虚拟环境

我个人喜欢用conda，管理起来方便。当然你用venv也行，看个人喜好。

# 创建Python 3.10环境
conda create -n yolov11_prune python=3.10 -y

# 激活环境
conda activate yolov11_prune

为什么选3.10？嗯，这里要注意。PyTorch 2.x对3.10支持最好，3.11以上有些底层算子还没完全适配。我在项目中踩过这个坑，当时图新鲜装了3.11，结果torch.compile报了一堆错。

3.3 安装PyTorch

PyTorch的安装要看你的硬件。先确认CUDA版本：

nvidia-smi

看到CUDA Version那一行了吗？比如显示12.1，那就装对应的版本。

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

# 如果你只有CPU（不推荐，剪枝后的微调会很慢）
# pip install torch torchvision torchaudio

3.4 安装Ultralytics

Ultralytics现在更新很快，YOLOv11需要8.3.0以上版本。直接装：

pip install ultralytics>=8.3.0

装完验证一下：

python -c "from ultralytics import YOLO; print(YOLO.__version__)"

如果看到版本号，说明装好了。你想想看，这一步要是报错，多半是PyTorch版本不匹配。Ultralytics对PyTorch版本有要求，太老或太新都可能出问题。

3.5 安装剪枝核心库

torch-pruning是我比较常用的结构化剪枝库。它支持对卷积层、BN层做通道剪枝，正好适合YOLO系列。

pip install torch-pruning>=1.4.0

安装后可以跑个简单测试：

python -c "import torch_pruning as tp; print(tp.__version__)"

3.6 安装辅助库

剩下的辅助库一起装：

pip install numpy opencv-python matplotlib tensorboard onnx onnxruntime

这里我多说一句。onnx和onnxruntime不是必须的，但建议装上。剪枝后我们通常要导出ONNX做推理加速，提前装好省得后面手忙脚乱。

3.7 环境验证脚本

所有东西装完后，我习惯跑一个完整的验证脚本。确保每个库都能正常导入，版本也对得上。

import torch
import torch_pruning as tp
from ultralytics import YOLO
import numpy as np
import cv2
import onnx
import onnxruntime

print("PyTorch版本:", torch.__version__)
print("CUDA可用:", torch.cuda.is_available())
print("torch-pruning版本:", tp.__version__)
print("Ultralytics版本:", YOLO.__version__)
print("NumPy版本:", np.__version__)
print("OpenCV版本:", cv2.__version__)
print("ONNX版本:", onnx.__version__)
print("ONNX Runtime版本:", onnxruntime.__version__)

# 测试YOLOv11模型加载
model = YOLO("yolo11n.pt")
print("YOLOv11模型加载成功！")

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

3.8 整体流程一览

我把整个环境搭建的流程画了张图，方便你对照着操作。

3.9 常见问题与解决

环境搭建时可能会遇到几个典型问题，我列出来供你参考。

• CUDA不可用： 检查驱动版本，或者重装对应CUDA版本的PyTorch。我曾经在服务器上折腾了半天，最后发现是LD_LIBRARY_PATH没设对。
• Ultralytics导入报错： 多半是PyTorch版本太旧。升级到2.0以上基本能解决。
• torch-pruning安装失败： 它依赖C++编译，Windows上需要装Visual Studio Build Tools。Linux上一般没问题。
• ONNX导出报错： 检查onnx和onnxruntime版本是否匹配。我建议两个装同一个大版本。