如何利用4090算力进行深度学习入门：从环境配置到第一个模型训练

1次阅读

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

NVIDIA RTX 4090 是目前消费级显卡中的算力王者，拥有 24GB GDDR6X 显存和超过 16,000 个 CUDA 核心。相比前代显卡，它在深度学习任务中能带来 2 - 3 倍的性能提升。对于刚入门的新手来说，4090 的优势在于：

大显存容量：可以训练更大批次的图像或更复杂的模型
高计算吞吐：缩短模型训练时间，加快实验迭代速度
最新架构支持：完整支持 DLSS 3、Tensor Core 等 AI 加速技术

首先确保系统已安装最新显卡驱动。在 Ubuntu 系统上可以通过以下命令安装：

sudo apt install nvidia-driver-535

推荐使用 CUDA 12.1 版本，这是目前对 4090 支持最好的版本：

wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
sudo sh cuda_12.1.0_530.30.02_linux.run

安装完成后，将 CUDA 加入环境变量：

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

cuDNN 是 NVIDIA 提供的深度学习加速库。下载对应 CUDA 12.1 的版本后：

sudo tar -xzvf cudnn-linux-x86_64-8.9.0.131_cuda12-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

运行以下命令验证环境是否正确配置：

import torch
print(torch.cuda.is_available())  # 应输出 True
print(torch.cuda.get_device_name(0))  # 应显示 RTX 4090

在 4090 上的基准测试结果（ResNet50，批量大小 32）：

框架	每秒图像数	显存占用
TensorFlow	215	18GB
PyTorch	230	16GB

对于新手建议选择 PyTorch，因为：

更直观的 API 设计
动态计算图更适合调试
社区支持更活跃

下面是一个完整的 MNIST 分类示例：

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim

# 1. 数据准备
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5,), (0.5,))
])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

# 2. 模型定义
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(torch.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = Net().cuda()  # 将模型移到 GPU

# 3. 训练配置
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练循环
for epoch in range(5):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.cuda(), labels.cuda()  # 数据移到 GPU

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 100 == 99:
            print(f'Epoch: {epoch+1}, Batch: {i+1}, Loss: {running_loss/100:.3f}')
            running_loss = 0.0

利用 4090 的 Tensor Core 加速计算：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, labels in trainloader:
    inputs, labels = inputs.cuda(), labels.cuda()

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

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

当显存不足时可以累积多个小批次的梯度：

accumulation_steps = 4

for i, (inputs, labels) in enumerate(trainloader):
    inputs, labels = inputs.cuda(), labels.cuda()

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

    scaler.scale(loss).backward()

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

CUDA out of memory
减少批量大小
使用梯度累积
清理不必要的缓存：torch.cuda.empty_cache()
CUDA 版本不匹配
检查 PyTorch/TensorFlow 版本是否支持当前 CUDA
使用 conda install pytorch torchvision torchaudio cudatoolkit=12.1 -c pytorch 确保版本匹配
性能不如预期
确保数据加载使用num_workers>0
检查是否启用了混合精度
使用 nvtop 监控 GPU 利用率