PyTorch下多GPU并行训练不等长时序样本的实现方法

要实现多GPU同时训练不同的单样本（适配不等长时序数据），可以通过PyTorch的torch.nn.parallel模块手动实现单进程多GPU并行，核心思路是将模型复制到每个GPU，分配单个样本到对应GPU独立计算，最后汇总梯度更新参数。

1. 模型定义 #

首先实现你的LSTM+Linear模型，确保能处理单个时序样本（输入形状为(seq_len, input_dim)）：

import torch
import torch.nn as nn
from torch.nn import parallel

class SeqModel(nn.Module):
    def __init__(self, input_dim=36, hidden_dim=64, output_dim=1):
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        # 将单样本转为(1, seq_len, input_dim)适配LSTM的batch_first格式
        x = x.unsqueeze(0)
        _, (h_n, _) = self.lstm(x)
        # 取LSTM最后一层的隐藏状态做预测
        out = self.fc(h_n.squeeze(0))
        return out

2. 并行训练核心代码 #

初始化组件 #

# 获取可用GPU数量
device_count = torch.cuda.device_count()
if device_count == 0:
    raise RuntimeError("无可用GPU设备")

# 初始化主模型（默认放到第一个GPU）
model = SeqModel().cuda(0)
# 复制模型到所有GPU
models = parallel.replicate(model, list(range(device_count)))

# 定义优化器（基于主模型参数）
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 你的batch_size=1的DataLoader
dataloader = ...  # 替换为你的DataLoader实例
data_iter = iter(dataloader)

训练循环 #

epochs = 10
for epoch in range(epochs):
    model.train()
    total_loss = 0.0
    batch_count = 0

    while True:
        samples, targets = [], []
        try:
            # 收集与GPU数量相等的单样本和目标
            for _ in range(device_count):
                x, y = next(data_iter)
                samples.append(x)
                targets.append(y)
            batch_count += 1
        except StopIteration:
            # 处理剩余不足GPU数量的样本
            if samples:
                sub_devices = list(range(len(samples)))
                sub_models = parallel.replicate(model, sub_devices)
                inputs = parallel.scatter(samples, sub_devices)
                sub_targets = parallel.scatter(targets, sub_devices)
                outputs = parallel.parallel_apply(sub_models, inputs, devices=sub_devices)
                
                losses = []
                for out, y in zip(outputs, sub_targets):
                    loss = nn.MSELoss()(out, y.cuda(out.device.index))
                    losses.append(loss)
                
                for loss in losses:
                    loss.backward()
                
                # 汇总梯度到主模型
                for param_idx, param in enumerate(model.parameters()):
                    grads = [sub_models[i].parameters().__next__() for i in range(len(sub_models))]
                    param.grad = torch.stack([g.grad for g in grads]).sum(dim=0)
                
                optimizer.step()
                optimizer.zero_grad()
                total_loss += sum(loss.item() for loss in losses)
            break

        # 分配样本和目标到对应GPU
        inputs = parallel.scatter(samples, list(range(device_count)))
        targets = parallel.scatter(targets, list(range(device_count)))

        # 每个GPU独立前向传播
        outputs = parallel.parallel_apply(models, inputs, devices=list(range(device_count)))

        # 计算每个GPU的损失
        losses = []
        for out, y in zip(outputs, targets):
            loss = nn.MSELoss()(out, y.cuda(out.device.index))
            losses.append(loss)

        # 反向传播计算梯度
        for loss in losses:
            loss.backward()

        # 汇总所有GPU的梯度到主模型
        for param_idx, param in enumerate(model.parameters()):
            grads = [models[i].parameters().__next__() for i in range(device_count)]
            # 梯度求和（或取平均，根据需求调整）
            param.grad = torch.stack([g.grad for g in grads]).sum(dim=0)

        # 更新主模型参数
        optimizer.step()
        optimizer.zero_grad()

        # 将更新后的参数同步到所有GPU的模型
        models = parallel.replicate(model, list(range(device_count)))

        # 统计损失
        total_loss += sum(loss.item() for loss in losses)

    avg_loss = total_loss / (batch_count * device_count + len(samples))
    print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")

3. 关键注意事项 #

• 样本分配：每次收集的样本数量需匹配GPU数量，最后一批样本不足时单独处理，避免浪费GPU资源。
• 梯度汇总：因为每个GPU的模型是主模型的复制，反向传播后需手动将所有GPU的梯度汇总到主模型，再执行参数更新。
• 参数同步：每次参数更新后，需将主模型的新参数重新复制到所有GPU的模型，保证下一轮训练参数一致。
• 损失函数：根据任务类型替换损失函数（如分类任务用nn.CrossEntropyLoss）。

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