3. 训练环境搭建
环境搭建这事儿,说简单也简单,说坑也多。我见过不少同学,模型写好了,一跑就报错,最后发现是CUDA版本没对上。嗯,咱们今天就把这块彻底理清楚。
3.1 虚拟环境管理:别让依赖打架
先聊虚拟环境。你想想看,一个项目用PyTorch 1.10,另一个用2.0,装在一起肯定打架。我个人习惯,每个项目一个独立环境,干净利落。
3.1.1 Conda vs. venv
| 工具 | 适用场景 | 优点 | 缺点 |
|---|---|---|---|
| Conda | 数据科学、深度学习 | 可管理Python版本和非Python依赖 | 环境较大,安装稍慢 |
| venv | 纯Python项目 | 轻量、内置、速度快 | 不能管理Python版本 |
我个人更推荐Conda,尤其是做模型移植时,经常需要切换Python版本。Conda一条命令搞定,省心。
核心命令速查
# 创建环境
conda create -n model_deploy python=3.9
# 激活环境
conda activate model_deploy
# 退出环境
conda deactivate
# 删除环境
conda env remove -n model_deploy
小技巧:环境名别起太随意。我习惯用 项目名_py版本 的格式,比如 yolo_v8_py39,一看就明白。
3.2 GPU驱动与CUDA配置:踩坑重灾区
这块是环境搭建里最容易出问题的。我曾经在一台新服务器上折腾了整整一下午,最后发现是驱动版本太老,跟CUDA 11.8不兼容。
3.2.1 先确认硬件
跑之前,先看看你的显卡支不支持。用这条命令:
nvidia-smi
如果显示显卡信息,说明驱动已装。如果没有,先去NVIDIA官网下载对应驱动。
注意:笔记本用户别用Windows自动更新的驱动,经常阉割了CUDA支持。去NVIDIA官网手动选型号下载最稳。
3.2.2 CUDA版本匹配
这里有个关键点:PyTorch/TensorFlow对CUDA版本有要求。我整理了一张对照表:
| 框架版本 | 推荐CUDA | 推荐cuDNN |
|---|---|---|
| PyTorch 2.0+ | 11.8 / 12.1 | 8.9+ |
| PyTorch 1.13 | 11.7 / 11.8 | 8.6+ |
| TensorFlow 2.12+ | 11.8 | 8.6+ |
| TensorFlow 2.10 | 11.2 | 8.1 |
为什么会这样?说白了,框架底层调用了CUDA的API,版本不对就报错。我建议先定框架版本,再装对应CUDA。
3.2.3 安装步骤
- 装驱动:nvidia-smi能跑就行
- 装CUDA Toolkit:去NVIDIA官网下载runfile或deb包
- 装cuDNN:解压后把文件拷到CUDA目录
- 配置环境变量:把CUDA路径加到PATH和LD_LIBRARY_PATH
环境变量配置示例(.bashrc)
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
3.3 PyTorch安装实战
PyTorch的安装其实很简单,关键选对命令。我个人习惯用pip安装,因为conda有时会装CPU版本。
3.3.1 根据CUDA版本选命令
去PyTorch官网,选择你的配置,会生成对应的安装命令。比如:
# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# CPU only
pip install torch torchvision torchaudio
验证安装:装完后跑一下这段代码,确认GPU可用。
import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))
如果返回True和显卡型号,恭喜你,成了。
3.4 TensorFlow安装实战
TensorFlow的安装稍微讲究一点。2.10是最后一个原生支持GPU的版本,之后的版本需要自己装tensorflow-cuda-plugin。
3.4.1 安装命令
# TensorFlow 2.10 (推荐)
pip install tensorflow==2.10
# TensorFlow 2.12+
pip install tensorflow
pip install nvidia-tensorflow
避坑指南:我曾经在Windows上装TensorFlow 2.12,折腾半天发现GPU用不了。后来查文档才知道,Windows上TensorFlow 2.10之后就不原生支持GPU了。建议Windows用户用WSL2。
3.4.2 验证安装
import tensorflow as tf
print(tf.__version__)
print(tf.config.list_physical_devices('GPU'))
如果输出不为空,说明GPU被正确识别。
3.5 环境搭建知识体系
我把整个流程画了张图,方便你对照检查:
3.6 常见问题与解决方案
最后,我把这些年遇到的典型问题列出来,你遇到了可以直接查:
| 问题现象 | 可能原因 | 解决方案 |
|---|---|---|
| torch.cuda.is_available() 返回False | CUDA版本不匹配 / 驱动太老 | 检查nvidia-smi版本,重装对应CUDA |
| TensorFlow找不到GPU | tensorflow版本与CUDA不兼容 | 降级到2.10或使用nvidia-tensorflow |
| ImportError: libcudart.so 找不到 | 环境变量没配好 | 检查LD_LIBRARY_PATH是否包含CUDA路径 |
| conda安装超慢 | 默认源在国外 | 换清华源或中科大源 |
终极建议:如果你刚开始接触,直接用Docker镜像。NVIDIA官方提供了PyTorch和TensorFlow的容器,驱动装好就能跑,省去所有配置烦恼。命令就一行:
docker run --gpus all -it nvcr.io/nvidia/pytorch:23.12-py3
环境搭建这块,说白了就是细心。每一步验证一下,别急着往下走。我刚开始也踩过不少坑,后来养成习惯:装一步,测一步,稳得很。