你好，我是云哥。最近研究了一下pytorch-lightning，寫了兩篇博客記錄，這是其一，其二有更多的驚喜，敬請期待。????

公眾號后臺回復(fù)關(guān)鍵字：源碼，獲取本文源代碼！

pytorch-lightning 是建立在pytorch之上的高層次模型接口。

pytorch-lightning 之于 pytorch，就如同keras之于 tensorflow。

通過使用 pytorch-lightning，用戶無需編寫自定義訓(xùn)練循環(huán)就可以非常簡潔地在CPU、單GPU、多GPU、乃至多TPU上訓(xùn)練模型。

無需考慮模型和數(shù)據(jù)在cpu,cuda之間的移動，并且可以通過回調(diào)函數(shù)實現(xiàn)CheckPoint參數(shù)保存，實現(xiàn)斷點續(xù)訓(xùn)功能。

一般按照如下方式 安裝和 引入 pytorch-lightning 庫。

#安裝
pip install pytorch-lightning

#引入
import pytorch_lightning as pl 

顧名思義，它可以幫助我們漂亮(pl)地進(jìn)行深度學(xué)習(xí)研究。????

一，pytorch-lightning的設(shè)計哲學(xué)

pytorch-lightning 的核心設(shè)計哲學(xué)是將 深度學(xué)習(xí)項目中的 研究代碼(定義模型) 和 工程代碼 (訓(xùn)練模型) 相互分離。

用戶只需專注于研究代碼(pl.LightningModule)的實現(xiàn)，而工程代碼借助訓(xùn)練工具類(pl.Trainer)統(tǒng)一實現(xiàn)。

更詳細(xì)地說，深度學(xué)習(xí)項目代碼可以分成如下4部分：

• 研究代碼 (Research code)，用戶繼承LightningModule實現(xiàn)。
• 工程代碼 (Engineering code)，用戶無需關(guān)注通過調(diào)用Trainer實現(xiàn)。
• 非必要代碼 （Non-essential research code，logging, etc...），用戶通過調(diào)用Callbacks實現(xiàn)。
• 數(shù)據(jù) (Data)，用戶通過torch.utils.data.DataLoader實現(xiàn)。

二，pytorch-lightning使用范例

下面我們使用minist圖片分類問題為例，演示pytorch-lightning的最佳實踐。

1，準(zhǔn)備數(shù)據(jù)

import torch 
from torch import nn 

import torchvision 
from torchvision import transforms

transform = transforms.Compose([transforms.ToTensor()])

ds_train = torchvision.datasets.MNIST(root="./minist/",train=True,download=True,transform=transform)
ds_valid = torchvision.datasets.MNIST(root="./minist/",train=False,download=True,transform=transform)

dl_train =  torch.utils.data.DataLoader(ds_train, batch_size=128, shuffle=True, num_workers=4)
dl_valid =  torch.utils.data.DataLoader(ds_valid, batch_size=128, shuffle=False, num_workers=4)

print(len(ds_train))
print(len(ds_valid))

Done!
60000
10000

2，定義模型

import pytorch_lightning as pl 
import datetime

class Model(pl.LightningModule):
    
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Conv2d(in_channels=1,out_channels=32,kernel_size = 3),
            nn.MaxPool2d(kernel_size = 2,stride = 2),
            nn.Conv2d(in_channels=32,out_channels=64,kernel_size = 5),
            nn.MaxPool2d(kernel_size = 2,stride = 2),
            nn.Dropout2d(p = 0.1),
            nn.AdaptiveMaxPool2d((1,1)),
            nn.Flatten(),
            nn.Linear(64,32),
            nn.ReLU(),
            nn.Linear(32,10)]
        )
        
    def forward(self,x):
        for layer in self.layers:
            x = layer(x)
        return x
    
    #定義loss,以及可選的各種metrics
    def training_step(self, batch, batch_idx):
        x, y = batch
        prediction = self(x)
        loss = nn.CrossEntropyLoss()(prediction,y)
        return loss
    
    #定義optimizer,以及可選的lr_scheduler
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return {"optimizer":optimizer}
    
    def validation_step(self, batch, batch_idx):
        loss = self.training_step(batch,batch_idx)
        return {"val_loss":loss}
    
    def test_step(self, batch, batch_idx):
        loss = self.training_step(batch,batch_idx)
        return {"test_loss":loss}
    
  

3，訓(xùn)練模型

pl.seed_everything(1234)
model = Model() 


ckpt_callback = pl.callbacks.ModelCheckpoint(
    monitor='val_loss',
    save_top_k=1,
    mode='min'
)

# gpus=0 則使用cpu訓(xùn)練，gpus=1則使用1個gpu訓(xùn)練，gpus=2則使用2個gpu訓(xùn)練，gpus=-1則使用所有g(shù)pu訓(xùn)練，
# gpus=[0,1]則指定使用0號和1號gpu訓(xùn)練， gpus="0,1,2,3"則使用0,1,2,3號gpu訓(xùn)練
# tpus=1 則使用1個tpu訓(xùn)練

trainer = pl.Trainer(max_epochs=5,gpus=0,callbacks = [ckpt_callback]) 

#斷點續(xù)訓(xùn)
#trainer = pl.Trainer(resume_from_checkpoint='./lightning_logs/version_31/checkpoints/epoch=02-val_loss=0.05.ckpt')

trainer.fit(model,dl_train,dl_valid)

Global seed set to 1234
GPU available: False, used: False
TPU available: None, using: 0 TPU cores

  | Name   | Type       | Params
--------------------------------------
0 | layers | ModuleList | 54.0 K
--------------------------------------
54.0 K    Trainable params
0         Non-trainable params
54.0 K    Total params
Epoch 4: 100% >>>>>>>>>>>>>>>>>>>>>>>>>>>> 158/158 [00:19<00:00, 8.08it/s, loss=0.138, v_num=34]

4，評估模型

result = trainer.test(model, test_dataloaders=dl_valid)
print(result)

--------------------------------------------------------------------------------
DATALOADER:0 TEST RESULTS
{'test_loss': tensor(0.0047)}
--------------------------------------------------------------------------------
[{'test_loss': 0.004680501762777567}]

5，使用模型

data,label = next(iter(dl_valid))
model.eval()
prediction = model(data)
print(prediction)

tensor([[ -5.1149,  -6.1142,   2.0591,  ...,   7.0609,  -5.4144,   0.5222],
        [ -2.2989,  -5.6076,   3.7343,  ...,  -1.8391,  -6.4941,  -3.4076],
        [  0.9215,   6.9357,  -1.9887,  ...,  -2.2996,  -0.8034,  -3.2993],
        ...,
        [ -4.5674,  -6.0223,  -0.9309,  ...,  -3.5468,   0.3367,   4.5473],
        [  4.3023,  -4.1629,  -1.2742,  ...,  -4.2527,  -2.3449,  -2.5585],
        [ -3.8913, -10.3790,  -1.7804,  ...,  -4.6757,  -0.7428,   1.0305]],
       grad_fn=
     
      )
      

     

6，保存模型

最優(yōu)模型默認(rèn)保存在 trainer.checkpoint_callback.best_model_path 的目錄下，可以直接加載。

print(trainer.checkpoint_callback.best_model_path)
print(trainer.checkpoint_callback.best_model_score)

/Users/liangyun/CodeFiles/PythonAiRoad/lightning_logs/version_34/checkpoints/epoch=04-val_loss=0.00.ckpt
tensor(0.0047)

model_clone = Model.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)
trainer_clone = pl.Trainer(max_epochs=3) 
result = trainer_clone.test(model_clone,dl_valid)
print(result)

--------------------------------------------------------------------------------
DATALOADER:0 TEST RESULTS
{'test_loss': tensor(0.0047)}
--------------------------------------------------------------------------------
[{'test_loss': 0.004680501762777567}]

如果對本文內(nèi)容理解上有需要進(jìn)一步和作者交流的地方，歡迎在公眾號"算法美食屋"下留言。作者時間和精力有限，會酌情予以回復(fù)。

也可以在公眾號后臺回復(fù)關(guān)鍵字：加群，加入讀者交流群和大家討論。