PyTorch Lightning 在validation_epoch_end 中将张量移动到正确的设备

Question

我想用 a 的validation_epoch_end方法创建一个新的张量LightningModule。从官方文档（第 48 页）中可以看出，我们应该避免直接.cuda()或.to(device)调用：

没有 .cuda() 或 .to() 调用。. . 闪电为你做这些。

我们鼓励使用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->设备传输

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

我有一个类似的问题来创建张量，这对我有帮助。我希望它也能帮助你。