【PyTorch Ligntning】如何将 PyTorch 组织为 Lightning

目录

一、移动运算代码 (computational code)

二、移动优化器 (optimizer(s)) 和调度器 (schedulers)

三、找到训练循环的主要部分

四、找到验证循环的主要部分

五、找到测试循环的主要部分

六、移除任何 .cuda() 或 to.device() 调用

项目地址：https://github.com/PyTorchLightning/pytorch-lightning

为使 PyTorch 代码能够用 Lightning 运行，本文将简要介绍如何 将 PyTorch 组织到 Lightning 中。

一、移动运算代码 (computational code)

将模型架构和前向传播移动到 LightningModule：

class LitModel(LightningModule):

    def __init__(self):
        super().__init__()
        self.layer_1 = torch.nn.Linear(28 * 28, 128)
        self.layer_2 = torch.nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        return x

二、移动优化器 (optimizer(s)) 和调度器 (schedulers)

将优化器移至 configure_optimizers() 钩子函数：

class LitModel(LightningModule):

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

三、找到训练循环的主要部分

Lightning 将自动执行大部分训练，epoch 和 batch 迭代，此处需要保留的只是训练步骤逻辑。这应进入 training_step() 钩子函数 (在这种情况下，请确保使用 hook parameters, batch 和 batch_idx)：

class LitModel(LightningModule):

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

四、找到验证循环的主要部分

要添加一个 (可选的) 验证循环，请向 validation_step() 钩子函数添加逻辑 (在这种情况下，请确保使用 hook parameters, batch 和 batch_idx)：

class LitModel(LightningModule):

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        val_loss = F.cross_entropy(y_hat, y)
        return val_loss

注意

model.eval() 和 torch.no_grad() 被自动调用用于验证

五、找到测试循环的主要部分

要添加一个 (可选的) 测试循环，请向 test_step() 钩子函数添加逻辑 (在这种情况下，请确保使用 hook parameters, batch 和 batch_idx)：

class LitModel(LightningModule):

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

注意

model.eval() 和 torch.no_grad() 被自动调用用于测试

如果没有主动调用，测试循环不会运行：

trainer.test()

小贴士

.test() 会自动加载最佳检查点 (checkpoint)

六、移除任何 .cuda() 或 to.device() 调用

因为你的 LightningModule 能够自动在任意硬件上运行！

参考文献

https://pytorch-lightning.readthedocs.io/en/stable/converting.html