如何在Pytorch中使用多显卡进行不等长样本训练？

Question

这是一个 Pytorch 问题。 我有一些样本，它们的形状像 (128, 36), (64, 36), (77, 36), ... 像这样，它们是一些长度不等的时间序列。 我的神经网络只包括 nn.LSTM 和 nn.Linear，我确信它可以处理不同长度的时间序列样本。 我有多个显卡。如何编写代码让每个显卡可以同时训练不同的单个样本？

我发现我只能创建一个batch_size为1的dataloader，然后使用单个显卡进行训练，即使torch.cuda.device_count() > 1。这太慢了。

Answer 1

为了在不等长度的样本上训练模型，但当 PyTorch 中有多个 GPU 时，您可能需要考虑以下方法：

自定义数据加载器 您可以编写一个自定义的数据加载器，每次仅返回一个序列。将batch_size设置为1，将确保数据加载器的每次迭代只会为您提供一个序列。

数据并行性 利用 DataParallel 或 DistributedDataParallel，您可以在多个 GPU 之间分配工作负载。由于您希望以单独的方式处理每个序列，因此您需要确保每个序列都进入不同的 GPU 进行处理。

手动分配 GPU 另一种方法是，您可以手动将每个样本分配给不同的 GPU。此后，您将返回结果。

让我通过一个简单的例子来详细说明：

模型定义：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleModel, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        h_0 = torch.zeros(1, x.size(0), hidden_dim).cuda(x.device.index)
        c_0 = torch.zeros(1, x.size(0), hidden_dim).cuda(x.device.index)
        out, _ = self.lstm(x, (h_0, c_0))
        out = self.fc(out[:, -1, :])
        return out

数据模拟：

# Assume you have 2 GPUs
num_gpus = 4
data_on_gpus = []

# Generate some data with variable sequence lengths
for _ in range(num_gpus):
    seq_len = torch.randint(60, 130, (1,)).item()
    data = torch.randn(1, seq_len, 36).cuda(_)
    data_on_gpus.append(data)

循环训练：

input_dim = 36
hidden_dim = 50
output_dim = 10

# Initialize the models on each GPU
models = [SimpleModel(input_dim, hidden_dim, output_dim).cuda(i) for i in range(num_gpus)]

# Initialize the optimizers
optimizers = [torch.optim.Adam(models[i].parameters()) for i in range(num_gpus)]

# Loop for training
for epoch in range(10):
    for i in range(num_gpus):
        optimizers[i].zero_grad()
        data = data_on_gpus[i]
        model = models[i]
        output = model(data)
        loss = torch.mean(output)
        loss.backward()
        optimizers[i].step()

我希望您觉得这有帮助。