При запуске LSTM с несколькими графическими процессорами появляется "Входные и скрытые тензоры не на одном устройстве"

Я пытаюсь тренировать слой LSTM в pytorch. Я использую 4 графических процессора. При инициализации я добавил функцию .cuda() для перемещения скрытого слоя в графический процессор. Но когда я запускаю код с несколькими графическими процессорами, я получаю эту ошибку во время выполнения:

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

Я попытался решить проблему с помощью функции .cuda() в функции forward, как показано ниже:

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

Эта линия, кажется, решает проблему, но вызывает у меня беспокойство, если обновленный скрытый слой виден в разных графических процессорах. Должен ли я переместить вектор обратно в процессор в конце функции пересылки для пакета или есть какой-либо другой способ решения проблемы.

ОТВЕТЫ

Ответ 1

Когда вы вызываете .cuda() для тензора, Pytorch перемещает его на текущее устройство с графическим процессором по умолчанию (GPU-0). Таким образом, из-за параллелизма данных ваши данные живут в другом графическом процессоре, а ваша модель переходит в другую, что приводит к ошибке времени выполнения, с которой вы сталкиваетесь.

Правильный способ реализации параллелизма данных для повторяющихся нейронных сетей заключается в следующем:

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)

Вам также необходимо установить переменную среды CUDA_VISIBLE_DEVICES соответствующим образом для настройки нескольких графических процессоров.

Ссылки: