PyTorch Lightning перемещает тензор для исправления устройства в validation_epoch_end

Question

Я хотел бы создать новый тензор в методе validation_epoch_end из LightningModule. В официальной документации (стр. 48) сказано, что нам следует избегать прямых вызовов .cuda() или .to(device):

.cuda () или .to () звонки. . . Lightning сделает это за вас.

, и мы рекомендуем использовать метод type_as для передачи на нужное устройство.

new_x = new_x.type_as(x.type())

Однако в шаг validation_epoch_end У меня нет никакого тензора для копирования устройства (методом type_as) чистым способом.

Мой вопрос: что мне делать, если я хочу создать новый тензор в этом метод и передать его на устройство, где находится модель?

Единственное, что я могу придумать, - это найти тензор в словаре outputs, но это выглядит как-то беспорядочно:

avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
output = self(self.__test_input.type_as(avg_loss))

Есть ли какой-нибудь чистый способ добиться этого?

Answer 1

вы проверили часть 3.4 (стр. 34) в do c, которую вы связали?

LightningModules знает, на каком устройстве они установлены! создавать тензоры на устройстве напрямую, чтобы избежать передачи CPU-> Device

t = tensor.rand(2, 2).cuda()# bad
(self is lightningModule)t = tensor.rand(2,2, device=self.device)# good 

У меня была аналогичная проблема для создания тензоров, это помогло мне. Надеюсь, это поможет и вам.