3. 训练环境搭建：PyTorch/TensorFlow安装、GPU驱动与CUDA配置、虚拟环境管理
环境搭建这事儿，说简单也简单，说坑也多。我见过不少同学，模型写好了，一跑就报错，最后发现是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_PATH3.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环境搭建这块，说白了就是细心。每一步验证一下，别急着往下走。我刚开始也踩过不少坑，后来养成习惯：装一步，测一步，稳得很。