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

Question

我想在 validation_epoch_end 的 LightningModule 方法中创建一个新的张量。从官方 docs(第 48 页)中指出，我们应该避免直接 .cuda() 或 .to(device) 调用:

There are no .cuda() or .to() calls. . . Lightning does these for you.

我们鼓励使用 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 know what device they are on! construct tensors on the device directly to avoid CPU->Device transfer

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

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