如何充分利用4090显卡算力：从环境配置到深度学习实战指南

1次阅读

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

刚拿到 RTX 4090 时，我也遇到过同样的问题——明明参数碾压旧显卡，实际训练速度却提升有限。经过两周踩坑总结，发现关键在于正确配置和优化。本文将手把手带你解锁 4090 的真实性能。

先到 NVIDIA 官网下载最新驱动（目前推荐 535+ 版本），安装后执行：

nvidia-smi  # 确认驱动版本和显卡识别正常

接着安装 CUDA Toolkit 12.1（与 4090 适配最佳）：

下载官方.run 安装包
执行安装时务必取消勾选自带的驱动安装
将以下内容添加到~/.bashrc：

export PATH=/usr/local/cuda-12.1/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH

去 NVIDIA 开发者网站下载与 CUDA 12.1 匹配的 cuDNN 8.9，解压后执行：

sudo cp cuda/include/cudnn*.h /usr/local/cuda/include/
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

PyTorch 用户推荐使用 2.0+ 版本：

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

TensorFlow 用户需注意：

pip install tensorflow[and-cuda]==2.12.0  # 最后一个官方支持版本

4090 的 Tensor Core 特别适合混合精度计算：

import torch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

效果：ResNet50 训练速度提升 2 - 3 倍，显存占用减少 30%

当遇到超大 batch 需求时：

accum_steps = 4  # 累积 4 个 batch 再更新

for i, (inputs, labels) in enumerate(train_loader):
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels) / accum_steps

    scaler.scale(loss).backward()

    if (i+1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

适用于迭代模式固定的场景：

g = torch.cuda.CUDAGraph()

# 首次运行捕获计算图
with torch.cuda.graph(g):
    static_output = model(static_input)
    static_loss = criterion(static_output, static_label)
    static_loss.backward()

# 后续训练直接复现
for input, label in dataloader:
    static_input.copy_(input)
    static_label.copy_(label)
    g.replay()  # 比常规执行快 15%

import torch
import time

device = 'cuda'
size = 8192  # 测试矩阵尺寸

a = torch.randn(size, size, device=device)
b = torch.randn(size, size, device=device)

# 预热 GPU
for _ in range(10):
    _ = torch.mm(a, b)

torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    c = torch.mm(a, b)
torch.cuda.synchronize()
print(f"TFLOPS: {2*size**3*100/(time.time()-start)/1e12:.2f}")

model = torchvision.models.resnet50().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(100):
    model.train()
    for inputs, labels in train_loader:
        inputs, labels = inputs.cuda(), labels.cuda()

        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

现象：CUDA out of memory
解决方案：
- 减小 batch size
- 使用梯度检查点（torch.utils.checkpoint）
- 清理无用变量（del tensor + torch.cuda.empty_cache()）
现象：训练中途突然崩溃
检查是否使用了 pin_memory=True 导致主机内存不足
现象：显存占用持续增加
确认没有在循环中持续累积计算图（with torch.no_grad()）