2. 环境搭建:Anaconda、PyTorch、CUDA、cuDNN 安装与验证
说实话,工业缺陷检测这个方向,环境搭不好,后面全是坑。
我见过太多人,代码写得挺溜,一跑就报错。最后发现,CUDA 版本不对,或者 PyTorch 和 cuDNN 没配上。白白浪费一整天。
这一章,咱们就把这事彻底搞定。我的原则是:一次配好,用到毕业。
2.1 整体架构:先看一张图
在动手之前,我建议你先搞清楚这些组件之间的关系。说白了,就是一层叠一层:
你看,从硬件到代码,每一层都得对上。尤其是 CUDA 版本,PyTorch 和 cuDNN 都依赖它。我习惯先查 PyTorch 官网,看它支持哪个 CUDA,再倒着装。
2.2 第一步:安装 Anaconda
Anaconda 这东西,说白了就是个 Python 环境管家。我强烈建议你用 Miniconda,轻量、干净,不装一堆你用不上的包。
2.2.1 下载与安装
- 去 清华镜像站 下载 Miniconda(别去官网,慢得要命)。
- 选 Python 3.10 或 3.11 版本。
- 安装时,勾选 "Add to PATH",省得后面麻烦。
C:\miniconda3 或 D:\miniconda3。
2.2.2 配置国内镜像源
这一步很多人忽略,结果装包慢到怀疑人生。直接上命令:
# 添加清华镜像源
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --set show_channel_urls yes
# 设置 pip 镜像(这个也很重要)
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
2.2.3 创建虚拟环境
我每个项目都建独立环境,互不干扰。工业项目尤其如此,依赖冲突太常见了。
# 创建环境,指定 Python 3.10
conda create -n yolo_env python=3.10 -y
# 激活环境
conda activate yolo_env
conda activate yolo_env。我曾经因为忘了激活环境,在 base 里装了一堆包,最后环境乱成一锅粥。
2.3 第二步:安装 NVIDIA 驱动与 CUDA
GPU 加速是 YOLO 训练的关键。没有 CUDA,你的显卡就是个摆设。
2.3.1 检查显卡驱动
先看看你的驱动版本够不够新:
# Windows 下
nvidia-smi
这个命令会显示驱动版本和最高支持的 CUDA 版本。比如显示 CUDA Version: 12.4,说明你的驱动最高能支持 CUDA 12.4。
2.3.2 安装 CUDA Toolkit
我推荐安装 CUDA 11.8 或 CUDA 12.1。这两个版本 PyTorch 官方支持最好。
- 去 NVIDIA 官网下载 CUDA Toolkit。
- 选 exe (local) 版本,安装更稳定。
- 安装时选 自定义,只勾选 "CUDA" 组件,别装全家桶。
装完后验证:
nvcc --version
看到版本号,说明 CUDA 装好了。
2.3.3 安装 cuDNN
cuDNN 是 CUDA 的加速库,专门优化深度学习计算。说白了,就是让卷积跑得更快。
- 去 NVIDIA 官网下载 cuDNN(需要注册账号)。
- 选对应 CUDA 11.x 或 12.x 的版本。
- 解压后,把
bin、include、lib三个文件夹复制到 CUDA 安装目录(默认C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1)。
2.4 第三步:安装 PyTorch
这一步最讲究版本匹配。我的原则是:去 PyTorch 官网,选对 CUDA 版本,复制命令直接装。
2.4.1 选择正确的安装命令
假设你装了 CUDA 11.8:
# CUDA 11.8 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
如果装了 CUDA 12.1:
# CUDA 12.1 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
2.4.2 验证 PyTorch 和 CUDA
装完后,跑一段验证代码:
import torch
print("PyTorch 版本:", torch.__version__)
print("CUDA 是否可用:", torch.cuda.is_available())
print("CUDA 版本:", torch.version.cuda)
print("cuDNN 版本:", torch.backends.cudnn.version())
print("GPU 数量:", torch.cuda.device_count())
print("GPU 名称:", torch.cuda.get_device_name(0))
如果输出 CUDA 是否可用: True,恭喜你,环境搭好了。
torch.cuda.is_available() 返回 False,99% 是 PyTorch 版本和 CUDA 版本不匹配。卸载重装,别纠结。
2.5 版本对照表
为了方便你对照,我整理了一张表。这是我多次踩坑后总结出来的:
| 组件 | 推荐版本 | 说明 |
|---|---|---|
| Python | 3.10 | 兼容性最好,YOLOv8 官方推荐 |
| CUDA Toolkit | 11.8 或 12.1 | PyTorch 官方支持最稳定 |
| cuDNN | 8.9.x (for CUDA 11.x) 9.x (for CUDA 12.x) |
必须匹配 CUDA 大版本 |
| PyTorch | 2.x (对应 CUDA 11.8/12.1) | 用 pip 安装,别用 conda |
| NVIDIA Driver | ≥ 525 (for CUDA 12.x) | 用 nvidia-smi 检查 |
2.6 完整验证流程
环境搭完,我习惯跑一个完整的验证脚本。确保所有组件都能协同工作:
import torch
import torchvision
import time
# 1. 检查版本
print(f"PyTorch: {torch.__version__}")
print(f"Torchvision: {torchvision.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
# 2. 创建测试张量
x = torch.randn(1000, 1000).cuda()
y = torch.randn(1000, 1000).cuda()
# 3. 跑一次矩阵乘法,验证 GPU 计算
start = time.time()
z = torch.mm(x, y)
torch.cuda.synchronize()
end = time.time()
print(f"GPU 矩阵乘法耗时: {end - start:.4f} 秒")
print("环境验证通过!")
2.7 常见问题与解决
- nvidia-smi 找不到:驱动没装好,去 NVIDIA 官网重装驱动。
- nvcc 找不到:CUDA 没加到 PATH。手动添加
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\bin到系统环境变量。 - torch.cuda.is_available() 返回 False:PyTorch 版本和 CUDA 不匹配。卸载 PyTorch,用官网命令重装。
- cuDNN 加载失败:cuDNN 文件没放对位置。检查
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\bin下有没有cudnn64_*.dll。
好了,环境搭好了。下一件事,就是让 YOLO 跑起来,看看我们的缺陷检测模型到底有多强。