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()
具体实现都通过类封装之后,主函数就显得简洁多了。