学习笔记——GPU训练

碱化钾2025-10-16 7:03

1.单机单卡配置

单机单卡配置指在一台计算机上使用单个GPU进行模型训练。这种配置适合小规模模型和数据集，配置简单，调试方便。

PyTorch实现代码：

python 复制代码

import torch
model = MyModel().to(device)  # 设备转移
optimizer = optim.Adam(model.parameters())
for data, targets in train_loader:
    data, targets = data.to(device), targets.to(device)
    outputs = model(data)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

2.单机多卡配置

单机多卡配置利用同一台计算机上的多个GPU协同训练，通过数据并行提高训练效率。适合中等规模模型训练。

PyTorch DDP实现：

python 复制代码

def train_ddp(rank, world_size):
    setup(rank, world_size)
    model = DDP(model_class().to(rank), device_ids=[rank])
    for epoch in range(epochs):
        for data, targets in train_loader:
            data, targets = data.to(rank), targets.to(rank)
            loss = criterion(model(data), targets)
            loss.backward()
            optimizer.step()
    cleanup()

3.多机多卡配置

多机多卡配置跨越多个计算节点，每个节点包含多个GPU，适合大规模模型训练。需要处理更复杂的通信和同步问题。

PyTorch多节点实现：

python 复制代码

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
model = DDP(model, device_ids=[local_rank])

4.性能优化建议

合理设置批量大小
优化数据加载管道
选择合适的通信后端
监控GPU利用率
调整学习率策略

5.适用场景选择

小规模实验：单机单卡
中等规模训练：单机多卡
大规模生产训练：多机多卡

每种配置都有其特点和适用场景，选择时需要综合考虑模型规模、数据集大小、硬件资源等因素。