使用多个 GPU 运行 LSTM 会得到“输入和隐藏张量不在同一设备上”

Posted 2023-03-12

【中文标题】使用多个 GPU 运行 LSTM 会得到“输入和隐藏张量不在同一设备上”

【英文标题】：Running LSTM with multiple GPUs gets "Input and hidden tensors are not at the same device"

【发布时间】：2019-06-27 22:27:10

【问题描述】：

我正在尝试在 pytorch 中训练 LSTM 层。我正在使用 4 个 GPU。初始化时，我添加了 .cuda() 函数将隐藏层移动到 GPU。但是当我使用多个 GPU 运行代码时，我得到了这个运行时错误：

RuntimeError: Input and hidden tensors are not at the same device

我尝试通过在 forward 函数中使用 .cuda() 函数来解决这个问题，如下所示：

self.hidden = (self.hidden[0].type(torch.FloatTensor).cuda(), self.hidden[1].type(torch.FloatTensor).cuda()) 

这条线似乎解决了这个问题，但它引起了我的担忧，即如果在不同的 GPU 中看到更新的隐藏层。我应该在批处理的前向函数结束时将向量移回 cpu，还是有其他方法可以解决问题。

【问题讨论】：

一开始你是如何在多个 GPU 上运行它的？你在用DataParallel吗？

是的，在初始化模型后，我正在运行这一行：model = torch.nn.DataParallel(model) 并且在运行代码时我正在设置 CUDA_VISIBLE_DEVICES

这个问题你解决了吗？谢谢

【参考方案1】：

当您在张量上调用 .cuda() 时，Pytorch 默认将其移动到 current GPU device (GPU-0)。因此，由于数据并行性，您的数据位于不同的 GPU 中，而您的模型转到另一个 GPU，这会导致您面临运行时错误。

递归神经网络实现数据并行的正确方法如下：

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class MyModule(nn.Module):
    # ... __init__, other methods, etc.

    # padded_input is of shape [B x T x *] (batch_first mode) and contains
    # the sequences sorted by lengths
    #   B is the batch size
    #   T is max sequence length
    def forward(self, padded_input, input_lengths):
        total_length = padded_input.size(1)  # get the max sequence length
        packed_input = pack_padded_sequence(padded_input, input_lengths,
                                            batch_first=True)
        packed_output, _ = self.my_lstm(packed_input)
        output, _ = pad_packed_sequence(packed_output, batch_first=True,
                                        total_length=total_length)
        return output

m = MyModule().cuda()
dp_m = nn.DataParallel(m)

您还需要为多 GPU 设置相应地设置 CUDA_VISIBLE_DEVICES 环境变量。

参考资料：

Data Parallelism Fast.ai Forums RNNs and Data Parallelism