Anaconda环境下高效调用GPU训练模型的避坑指南

1次阅读

共计 3811 个字符，预计需要花费 10 分钟才能阅读完成。

在 Anaconda 中使用 GPU 训练深度学习模型时，开发者常常遇到以下问题：

CUDA 版本不匹配：安装的深度学习框架版本与本地 CUDA 驱动版本不兼容，导致无法调用 GPU
环境冲突：多个项目使用不同版本的 CUDA/cuDNN，导致环境混乱
GPU 显存不足：未合理设置 batch size 导致显存溢出，训练过程中断
安装方式混乱：conda 和 pip 混用导致依赖冲突
GPU 利用率低：未正确配置导致 GPU 计算资源未被充分利用

conda 安装
优点：自动解决 CUDA/cuDNN 依赖，环境隔离性好
缺点：版本更新可能滞后
pip 安装
优点：版本更新及时
缺点：需要手动管理 CUDA/cuDNN 依赖

推荐使用 conda 创建独立环境并安装 GPU 版本框架，可避免大多数兼容性问题。

框架版本	推荐 CUDA 版本	推荐 cuDNN 版本
PyTorch 1.12	11.3	8.2
TensorFlow 2.9	11.2	8.1

创建新 conda 环境
```
conda create -n gpu_env python=3.8
```
激活环境
```
conda activate gpu_env
```

安装 GPU 版本框架

# PyTorch
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

# TensorFlow
conda install tensorflow-gpu=2.9 cudatoolkit=11.2 cudnn=8.1

import torch

print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 是否可用: {torch.cuda.is_available()}")
print(f"GPU 数量: {torch.cuda.device_count()}")
print(f"当前 GPU: {torch.cuda.current_device()}")
print(f"GPU 名称: {torch.cuda.get_device_name(0)}")

import tensorflow as tf

print(f"TensorFlow 版本: {tf.__version__}")
print(f"GPU 是否可用: {tf.test.is_gpu_available()}")
print(f"GPU 设备列表: {tf.config.list_physical_devices('GPU')}")

命令行实时监控
```
watch -n 1 nvidia-smi
```

Python 代码监控

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
print(f"GPU 利用率: {util.gpu}%")
print(f"显存利用率: {util.memory}%")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 创建示例数据
X = torch.randn(1000, 10)
y = torch.randn(1000, 1)
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 定义简单模型
model = nn.Sequential(nn.Linear(10, 50),
    nn.ReLU(),
    nn.Linear(50, 1)
).to(device)  # 将模型转移到 GPU

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
for epoch in range(10):
    for batch_X, batch_y in dataloader:
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)  # 数据转移到 GPU

        optimizer.zero_grad()
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

import tensorflow as tf

# 限制 GPU 内存增长，避免一次性占用所有显存
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# 创建模型时自动使用 GPU
model = tf.keras.Sequential([tf.keras.layers.Dense(50, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1)
])

model.compile(optimizer='adam', loss='mse')

# 训练数据会自动在 CPU 和 GPU 之间传输
model.fit(X.numpy(), y.numpy(), epochs=10, batch_size=32)

增大 batch size 可以提高 GPU 利用率，但会增加显存占用

建议使用以下公式估算最大 batch size：

最大 batch size = (总显存 - 模型参数占用显存) / 单个样本占用显存

可以通过 torch.cuda.empty_cache() 手动释放未使用的显存

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in range(10):
    for batch_X, batch_y in dataloader:
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)

        optimizer.zero_grad()

        # 启用混合精度
        with autocast():
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)

        # 缩放梯度
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

CUDA out of memory

解决方案：减小 batch size，使用梯度累积

# 梯度累积示例
accumulation_steps = 4
for i, (batch_X, batch_y) in enumerate(dataloader):
    loss = model(batch_X, batch_y) / accumulation_steps
    loss.backward()

    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()