0. 简介
pytorch lightning通过提供LightningModule和LightningDataModule，使得在用pytorch编写网络模型时，加载数据、分割数据集、训练、验证、测试、计算指标的代码全部都能很好的组织起来，显得主程序调用时，代码简洁可读性大幅度提升。

1. pytorch lightning的安装

1 pip install pytorch-lightning
2 conda install pytorch-lightning -c conda-forge

2. 定义一个网络模型模型：LightningModule

通过继承LightningModule，并实现几个关键的函数，使得模型在训练、验证和测试过程中能进行模块化调用，具体细节完全被自定义的函数封装，整体十分简洁。定义一个LightningModule的基类，可以实现的函数如下：

 1 from pytorch_lightning import LightningModule
 2  
 3 class MyModel(LightningModule):
 4     """
 5     The only required methods in the LightningModule are:
 6     init
 7     training_step
 8     configure_optimizers
 9     """
10     def __init__(self, *args, **kwargs): pass
11     def forward(self, *args, **kwargs): pass
12     def training_step(self, batch, batch_idx, optimizer_idx, hiddens): pass
13     def training_step_end(self, *args, **kwargs): pass # 接受train_step的返回值
14     def training_epoch_end(self, outputs): pass # 接受train_step一整个epoch的返回值的列表
15     def validation_step(self, batch, batch_idx, dataloader_idx): pass # model.eval() and torch.no_grad() are called automatically
16     def validation_step_end(self, *args, **kwargs): pass # 接受validation_step的返回值
17     def validation_epoch_end(self, outputs)
18     def test_step(self, batch, batch_idx, dataloader_idx): pass # model.eval() and torch.no_grad() are called automatically
19     def test_step_end(self, *args, **kwargs): pass  # 接收test_step的返回值
20     def test_epoch_end(self, outputs): pass
21     def configure_optimizers(self, ): pass
22     def any_extra_hook(...): pass  #  指代任意其他的可重载函数

其中，必须实现的函数只有__init__() 、training_step()、configure_optimizers()。

3. 定义一个数据模型：LightningDataModule

通过定义LightningDataModule的子类，数据集分割、加载的代码将整合在一起，可以实现的方法有：

 1 class MyDataModule(LightningDataModule):
 2     def __init__(self):
 3         super().__init__()
 4     def prepare_data(self):
 5         # download, split, etc...
 6         # only called on 1 GPU/TPU in distributed
 7     def setup(self,stage:str):  # stage: "fit", "test", 【暂时不知道验证步骤叫什么名字，可以自己打印一下】
 8         # make assignments here (val/train/test split)
 9         # called on every process in DDP
10     def train_dataloader(self):
11         train_split = Dataset(...)
12         return DataLoader(train_split)
13     def val_dataloader(self):
14         val_split = Dataset(...)
15         return DataLoader(val_split)
16     def test_dataloader(self):
17         test_split = Dataset(...)
18         return DataLoader(test_split)

4. 使用pytorch lightning的API开始训练

 1 def main():
 2     model = MyModule()
 3     data_module = MyDataModule()
 4     trainer = pytorch_lightning.Trainer(...)  # some arugments, 根据需要传入你的参数
 5     trainer.fit(module, datamodule=data_module)
 6     trainer.test(module, datamodule=data_module, verbose=True)
 7  
 8  
 9 if __name__ == "__main__":
10     main()

具体实现都通过类封装之后，主函数就显得简洁多了。

标签：__,简要,self,pytorch,lightning,step,pass,test,def
来源： https://www.cnblogs.com/ltkekeli1229/p/15732291.html