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) в документе, на который вы ссылаетесь?

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

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

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