如何在Kaggle上利用双T4 GPU进行PyTorch模型训练？

一、先确认GPU可用性 #

先运行以下代码，确认PyTorch能识别到两块T4：

import torch
print(f"可用GPU数量: {torch.cuda.device_count()}")
print(f"GPU型号: {torch.cuda.get_device_name(0)}, {torch.cuda.get_device_name(1)}")

如果输出显示2块T4，就可以继续下一步。

二、多GPU训练方式选择（推荐第一种） #

在Kaggle这种单节点多GPU场景下，DataParallel是最易用的方案，能快速适配现有代码；若追求极致效率可尝试DistributedDataParallel，但配置复杂度更高，这里先讲前者。

三、修改你的训练代码 #

1. 包装模型并移至GPU

初始化模型后，用torch.nn.DataParallel包装模型，自动调用所有可用GPU：

# 假设已完成model、criterion、optimizer的初始化
model = model.cuda()
model = torch.nn.DataParallel(model)
# 也可以手动指定GPU：model = torch.nn.DataParallel(model, device_ids=[0,1])

2. 将输入数据移至GPU

训练循环中，把inputs和labels同步到GPU上：

inputs, labels = inputs.cuda(), labels.cuda()

3. 修改后的完整训练代码

import torch
from tqdm import tqdm

# 假设已完成dataset、model、criterion、optimizer的初始化

# 包装模型到多GPU
model = model.cuda()
model = torch.nn.DataParallel(model)

for epoch in range(2):
    running_loss = 0.0
    for data in tqdm(dataset):
        inputs, labels = data
        # 将数据移至GPU
        inputs, labels = inputs.cuda(), labels.cuda()
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(dataset)}")

四、额外注意事项 #

• Batch Size调整：DataParallel会把一个大批次拆分成多个子批次分配到不同GPU，建议适当放大原batch size（比如从32调到64），充分利用双GPU显存。
• Kaggle环境校验：确保Notebook设置中已启用GPU加速（右上角「Settings」→「Accelerator」选择GPU）。
• 模型保存与加载：用DataParallel包装后的模型，保存时会带module.前缀，加载时若无需多GPU，可移除前缀：

  # 保存模型
  torch.save(model.state_dict(), "model.pth")
  # 加载模型（单GPU场景）
  model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load("model.pth").items()})
  

内容的提问来源于stack exchange，提问作者Ilya Rudenko