如何在Kaggle中使用PyTorch调用双GPU进行模型训练？

核心解决方案 #

PyTorch提供两种主流多GPU训练方案，适配Kaggle双GPU环境，直接套用即可：

方法1：使用nn.DataParallel（快速适配） #

无需大幅修改现有单GPU代码，适合中小型模型：

import torch
import torch.nn as nn

# 先确认可用GPU数量
print(f"可用GPU数量: {torch.cuda.device_count()}")

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = YourModelClass()  # 替换为你的模型类

# 当存在多个GPU时包装模型
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

model.to(device)

提示：训练时输入数据仍需传到device，和单GPU训练逻辑一致；若需访问原模型属性/方法，使用model.module。

方法2：使用nn.parallel.DistributedDataParallel（性能优先） #

适合大型模型，训练效率更高，Kaggle环境原生支持：

1. 初始化分布式训练环境：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

# 初始化进程组（NCCL是GPU分布式训练的最优后端）
dist.init_process_group(backend='nccl')
local_rank = dist.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device('cuda', local_rank)

# 加载模型并绑定到当前GPU
model = YourModelClass().to(device)
model = DDP(model, device_ids=[local_rank])

2. 配合DistributedSampler拆分数据集（避免多GPU重复加载数据）：

train_dataset = YourDataset()  # 替换为你的数据集类
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=your_batch_size,  # 单GPU批次大小，总批次为该值×GPU数量
    sampler=train_sampler
)

3. 训练循环中更新sampler的epoch（保证每轮数据拆分不同）：

for epoch in range(total_epochs):
    train_sampler.set_epoch(epoch)
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        # 常规训练逻辑（前向传播、计算损失、反向传播、优化器更新）

4. 训练结束后清理进程组：

dist.destroy_process_group()

关键注意事项 #

• 先通过torch.cuda.device_count()确认Kaggle环境是否真的分配了2块GPU，若返回1则需检查Kaggle的GPU设置
• 使用DataParallel时，模型输出会是各GPU的结果列表，可通过.mean()或torch.cat()统一处理
• 两种方案都要求输入数据与模型在同一GPU设备上，避免出现“张量与模型不在同一设备”的错误

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