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

Posted 2023-03-27

【中文标题】PyTorch Lightning 将张量移动到 validation_epoch_end 中的正确设备

【英文标题】：PyTorch Lightning move tensor to correct device in validation_epoch_end

【发布时间】：2020-10-29 04:22:52

【问题描述】：

我想在LightningModule 的validation_epoch_end 方法中创建一个新张量。来自官方docs（第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))

有什么干净的方法可以实现吗？

【参考方案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 

我在创建张量时遇到了类似的问题，这对我很有帮助。希望对你也有帮助。

【讨论】：

谢谢，这正是我要找的！

当我读到你的（工作）答案时，我正要哭了：谢谢！

文档有所更改，但代码仍然有效： device 模块所在的设备。使用它来保持您的代码设备不可知。 def training_step(self): z = torch.rand(2, 3, device=self.device)