pytorch lightning使用（簡要介紹）

0. 簡介
pytorch lightning通過提供LightningModule和LightningDataModule，使得在用pytorch編寫網絡模型時，加載數據、分割數據集、訓練、驗證、測試、計算指標的代碼全部都能很好的組織起來，顯得主程序調用時，代碼簡潔可讀性大幅度提升。

1. pytorch lightning的安裝

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

2. 定義一個網絡模型模型：LightningModule

通過繼承LightningModule，並實現幾個關鍵的函數，使得模型在訓練、驗證和測試過程中能進行模塊化調用，具體細節完全被自定義的函數封裝，整體十分簡潔。定義一個LightningModule的基類，可以實現的函數如下：

from pytorch_lightning import LightningModule
 
class MyModel(LightningModule):
    """
    The only required methods in the LightningModule are:
    init
    training_step
    configure_optimizers
    """
    def __init__(self, *args, **kwargs): pass
    def forward(self, *args, **kwargs): pass
    def training_step(self, batch, batch_idx, optimizer_idx, hiddens): pass
    def training_step_end(self, *args, **kwargs): pass # 接受train_step的返回值
    def training_epoch_end(self, outputs): pass # 接受train_step一整個epoch的返回值的列表
    def validation_step(self, batch, batch_idx, dataloader_idx): pass # model.eval() and torch.no_grad() are called automatically
    def validation_step_end(self, *args, **kwargs): pass # 接受validation_step的返回值
    def validation_epoch_end(self, outputs)
    def test_step(self, batch, batch_idx, dataloader_idx): pass # model.eval() and torch.no_grad() are called automatically
    def test_step_end(self, *args, **kwargs): pass  # 接收test_step的返回值
    def test_epoch_end(self, outputs): pass
    def configure_optimizers(self, ): pass
    def any_extra_hook(...): pass  #  指代任意其他的可重載函數

其中，必須實現的函數只有__init__() 、training_step()、configure_optimizers()。

3. 定義一個數據模型：LightningDataModule

通過定義LightningDataModule的子類，數據集分割、加載的代碼將整合在一起，可以實現的方法有：

class MyDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
    def prepare_data(self):
        # download, split, etc...
        # only called on 1 GPU/TPU in distributed
    def setup(self,stage:str):  # stage: "fit", "test", 【暫時不知道驗證步驟叫什么名字，可以自己打印一下】
        # make assignments here (val/train/test split)
        # called on every process in DDP
    def train_dataloader(self):
        train_split = Dataset(...)
        return DataLoader(train_split)
    def val_dataloader(self):
        val_split = Dataset(...)
        return DataLoader(val_split)
    def test_dataloader(self):
        test_split = Dataset(...)
        return DataLoader(test_split)

4. 使用pytorch lightning的API開始訓練

def main():
    model = MyModule()
    data_module = MyDataModule()
    trainer = pytorch_lightning.Trainer(...)  # some arugments, 根據需要傳入你的參數
    trainer.fit(module, datamodule=data_module)
    trainer.test(module, datamodule=data_module, verbose=True)
 
 
if __name__ == "__main__":
    main()

具體實現都通過類封裝之后，主函數就顯得簡潔多了。