• pytorch_lightning模型训练加速技巧与涨点技巧

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

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

    pytorch-lightning 有以下一些引人注目的功能：

    • 可以不必编写自定义循环，只要指定loss计算方法即可。

    • 可以通过callbacks非常方便地添加CheckPoint参数保存、early_stopping 等功能。

    • 可以非常方便地在单CPU、多CPU、单GPU、多GPU乃至多TPU上训练模型。

    • 可以通过调用torchmetrics库，非常方便地添加Accuracy,AUC,Precision等各种常用评估指标。

    • 可以非常方便地实施多批次梯度累加、半精度混合精度训练、最大batch_size自动搜索等技巧，加快训练过程。

    • 可以非常方便地使用SWA(随机参数平均)、CyclicLR(学习率周期性调度策略)与auto_lr_find(最优学习率发现)等技巧 实现模型涨点。

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

    1. #安装
    2. pip install pytorch-lightning
    1. #引入
    2. import pytorch_lightning as pl

    顾名思义，它可以帮助我们漂亮(pl)地进行深度学习研究。😋😋 You do the research. Lightning will do everything else.⭐️⭐️

    参考文档：

    • pl_docs: https://pytorch-lightning.readthedocs.io/en/latest/starter/introduction.html

    • pl_template:https://github.com/PyTorchLightning/deep-learning-project-template

    • torchmetrics: https://torchmetrics.readthedocs.io/en/latest/pages/lightning.html

    公众号后台回复关键词：pl，获取本文jupyter notebook源代码。

    一，pytorch-lightning的设计哲学

    pytorch-lightning 的核心设计哲学是将 深度学习项目中的 研究代码(定义模型) 和 工程代码 (训练模型) 相互分离。

    用户只需专注于研究代码(pl.LightningModule)的实现，而工程代码借助训练工具类(pl.Trainer)统一实现。

    更详细地说，深度学习项目代码可以分成如下4部分：

    • 研究代码 (Research code)，用户继承LightningModule实现。

    • 工程代码 (Engineering code)，用户无需关注通过调用Trainer实现。

    • 非必要代码 （Non-essential research code，logging, etc...），用户通过调用Callbacks实现。

    • 数据 (Data)，用户通过torch.utils.data.DataLoader实现，也可以封装成pl.LightningDataModule。

    二，pytorch-lightning使用范例

    下面我们使用minist图片分类问题为例，演示pytorch-lightning的最佳实践。

    1，准备数据

    1. import torch 
    2. from torch import nn 
    3. from torchvision import transforms as T
    4. from torchvision.datasets import MNIST
    5. from torch.utils.data import DataLoader,random_split
    6. import pytorch_lightning as pl 
    7. from torchmetrics import Accuracy
    1. class MNISTDataModule(pl.LightningDataModule):
    2.     def __init__(self, data_dir: str = "./minist/", 
    3.                  batch_size: int = 32,
    4.                  num_workers: int =4):
    5.         super().__init__()
    6.         self.data_dir = data_dir
    7.         self.batch_size = batch_size
    8.         self.num_workers = num_workers
    9.  
    10.     def setup(self, stage = None):
    11.         transform = T.Compose([T.ToTensor()])
    12.         self.ds_test = MNIST(self.data_dir, train=False,transform=transform,download=True)
    13.         self.ds_predict = MNIST(self.data_dir, train=False,transform=transform,download=True)
    14.         ds_full = MNIST(self.data_dir, train=True,transform=transform,download=True)
    15.         self.ds_train, self.ds_val = random_split(ds_full, [55000, 5000])
    16.  
    17.     def train_dataloader(self):
    18.         return DataLoader(self.ds_train, batch_size=self.batch_size,
    19.                           shuffle=True, num_workers=self.num_workers,
    20.                           pin_memory=True)
    21.  
    22.     def val_dataloader(self):
    23.         return DataLoader(self.ds_val, batch_size=self.batch_size,
    24.                           shuffle=False, num_workers=self.num_workers,
    25.                           pin_memory=True)
    26.  
    27.     def test_dataloader(self):
    28.         return DataLoader(self.ds_test, batch_size=self.batch_size,
    29.                           shuffle=False, num_workers=self.num_workers,
    30.                           pin_memory=True)
    31.  
    32.     def predict_dataloader(self):
    33.         return DataLoader(self.ds_predict, batch_size=self.batch_size,
    34.                           shuffle=False, num_workers=self.num_workers,
    35.                           pin_memory=True)
    1. data_mnist = MNISTDataModule()
    2. data_mnist.setup()
    1. for features,labels in data_mnist.train_dataloader():
    2.     print(features.shape)
    3.     print(labels.shape)
    4.     break
    1. torch.Size([32, 1, 28, 28])
    2. torch.Size([32])

    2，定义模型

    1. net = nn.Sequential(
    2.     nn.Conv2d(in_channels=1,out_channels=32,kernel_size = 3),
    3.     nn.MaxPool2d(kernel_size = 2,stride = 2),
    4.     nn.Conv2d(in_channels=32,out_channels=64,kernel_size = 5),
    5.     nn.MaxPool2d(kernel_size = 2,stride = 2),
    6.     nn.Dropout2d(p = 0.1),
    7.     nn.AdaptiveMaxPool2d((1,1)),
    8.     nn.Flatten(),
    9.     nn.Linear(64,32),
    10.     nn.ReLU(),
    11.     nn.Linear(32,10)
    12. )
    13.  
    14. class Model(pl.LightningModule):
    15.     
    16.     def __init__(self,net,learning_rate=1e-3):
    17.         super().__init__()
    18.         self.save_hyperparameters()
    19.         self.net = net
    20.         self.train_acc = Accuracy()
    21.         self.val_acc = Accuracy()
    22.         self.test_acc = Accuracy() 
    23.         
    24.         
    25.     def forward(self,x):
    26.         x = self.net(x)
    27.         return x
    28.     
    29.     
    30.     #定义loss
    31.     def training_step(self, batch, batch_idx):
    32.         x, y = batch
    33.         preds = self(x)
    34.         loss = nn.CrossEntropyLoss()(preds,y)
    35.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    36.     
    37.     #定义各种metrics
    38.     def training_step_end(self,outputs):
    39.         train_acc = self.train_acc(outputs['preds'], outputs['y']).item()    
    40.         self.log("train_acc",train_acc,prog_bar=True)
    41.         return {"loss":outputs["loss"].mean()}
    42.     
    43.     #定义optimizer,以及可选的lr_scheduler
    44.     def configure_optimizers(self):
    45.         return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
    46.     
    47.     def validation_step(self, batch, batch_idx):
    48.         x, y = batch
    49.         preds = self(x)
    50.         loss = nn.CrossEntropyLoss()(preds,y)
    51.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    52.  
    53.     def validation_step_end(self,outputs):
    54.         val_acc = self.val_acc(outputs['preds'], outputs['y']).item()    
    55.         self.log("val_loss",outputs["loss"].mean(),on_epoch=True,on_step=False)
    56.         self.log("val_acc",val_acc,prog_bar=True,on_epoch=True,on_step=False)
    57.     
    58.     def test_step(self, batch, batch_idx):
    59.         x, y = batch
    60.         preds = self(x)
    61.         loss = nn.CrossEntropyLoss()(preds,y)
    62.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    63.     
    64.     def test_step_end(self,outputs):
    65.         test_acc = self.test_acc(outputs['preds'], outputs['y']).item()    
    66.         self.log("test_acc",test_acc,on_epoch=True,on_step=False)
    67.         self.log("test_loss",outputs["loss"].mean(),on_epoch=True,on_step=False)
    68.     
    69. model = Model(net)
    70.  
    71. #查看模型大小
    72. model_size = pl.utilities.memory.get_model_size_mb(model)
    73. print("model_size = {} M \n".format(model_size))
    74. model.example_input_array = [features]
    75. summary = pl.utilities.model_summary.ModelSummary(model,max_depth=-1)
    76. print(summary)
    1. model_size = 0.218447 M 
    2.  
    3.    | Name      | Type              | Params | In sizes         | Out sizes       
    4. ---------------------------------------------------------------------------------------
    5. 0  | net       | Sequential        | 54.0 K | [32, 1, 28, 28]  | [32, 10]        
    6. 1  | net.0     | Conv2d            | 320    | [32, 1, 28, 28]  | [32, 32, 26, 26]
    7. 2  | net.1     | MaxPool2d         | 0      | [32, 32, 26, 26] | [32, 32, 13, 13]
    8. 3  | net.2     | Conv2d            | 51.3 K | [32, 32, 13, 13] | [32, 64, 9, 9]  
    9. 4  | net.3     | MaxPool2d         | 0      | [32, 64, 9, 9]   | [32, 64, 4, 4]  
    10. 5  | net.4     | Dropout2d         | 0      | [32, 64, 4, 4]   | [32, 64, 4, 4]  
    11. 6  | net.5     | AdaptiveMaxPool2d | 0      | [32, 64, 4, 4]   | [32, 64, 1, 1]  
    12. 7  | net.6     | Flatten           | 0      | [32, 64, 1, 1]   | [32, 64]        
    13. 8  | net.7     | Linear            | 2.1 K  | [32, 64]         | [32, 32]        
    14. 9  | net.8     | ReLU              | 0      | [32, 32]         | [32, 32]        
    15. 10 | net.9     | Linear            | 330    | [32, 32]         | [32, 10]        
    16. 11 | train_acc | Accuracy          | 0      | ?                | ?               
    17. 12 | val_acc   | Accuracy          | 0      | ?                | ?               
    18. 13 | test_acc  | Accuracy          | 0      | ?                | ?               
    19. ---------------------------------------------------------------------------------------
    20. 54.0 K    Trainable params
    21. 0         Non-trainable params
    22. 54.0 K    Total params
    23. 0.216     Total estimated model params size (MB）

    3，训练模型

    1. pl.seed_everything(1234)
    2.  
    3. ckpt_callback = pl.callbacks.ModelCheckpoint(
    4.     monitor='val_loss',
    5.     save_top_k=1,
    6.     mode='min'
    7. )
    8. early_stopping = pl.callbacks.EarlyStopping(monitor = 'val_loss',
    9.                patience=3,
    10.                mode = 'min')
    11.  
    12. # gpus=0 则使用cpu训练，gpus=1则使用1个gpu训练，gpus=2则使用2个gpu训练，gpus=-1则使用所有gpu训练，
    13. # gpus=[0,1]则指定使用0号和1号gpu训练， gpus="0,1,2,3"则使用0,1,2,3号gpu训练
    14. # tpus=1 则使用1个tpu训练
    15.  
    16. trainer = pl.Trainer(max_epochs=20,   
    17.      #gpus=0, #单CPU模式
    18.      gpus=0, #单GPU模式
    19.      #num_processes=4,strategy="ddp_find_unused_parameters_false", #多CPU(进程)模式
    20.      #gpus=[0,1,2,3],strategy="dp", #多GPU的DataParallel(速度提升效果一般)
    21.      #gpus=[0,1,2,3],strategy=“ddp_find_unused_parameters_false" #多GPU的DistributedDataParallel(速度提升效果好)
    22.      callbacks = [ckpt_callback,early_stopping],
    23.      profiler="simple") 
    25. #断点续训
    26. #trainer = pl.Trainer(resume_from_checkpoint='./lightning_logs/version_31/checkpoints/epoch=02-val_loss=0.05.ckpt')
    28. #训练模型
    29. trainer.fit(model,data_mnist)
    1. Epoch 8: 100%
    2. 1876/1876 [01:44<00:00, 17.93it/s, loss=0.0603, v_num=0, train_acc=1.000, val_acc=0.985]

    4，评估模型

    result = trainer.test(model,data_mnist.train_dataloader(),ckpt_path='best')
    1. --------------------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9966545701026917, 'test_loss': 0.010617421939969063}
    4. --------------------------------------------------------------------------------
    result = trainer.test(model,data_mnist.val_dataloader(),ckpt_path='best')
    1. --------------------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9865999817848206, 'test_loss': 0.042671505361795425}
    4. --------------------------------------------------------------------------------
    result = trainer.test(model,data_mnist.test_dataloader(),ckpt_path='best')
    1. --------------------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.987500011920929, 'test_loss': 0.047178059816360474}
    4. --------------------------------------------------------------------------------

    5，使用模型

    1. data,label = next(iter(data_module.test_dataloader()))
    2. model.eval()
    3. prediction = model(data)
    4. print(prediction)
    1. tensor([[-13.0112,  -2.8257,  -1.8588,  -3.6137,  -0.3307,  -5.4953, -19.7282,
    2.           15.9651,  -8.0379,  -2.2925],
    3.         [ -6.0261,  -2.5480,  13.4140,  -5.5701, -10.2049,  -6.4469,  -3.7119,
    4.           -6.0732,  -6.0826,  -7.7339],
    5.           ...
    6.         [-16.7028,  -4.9060,   0.4400,  24.4337, -12.8793,   1.5085, -17.9232,
    7.           -3.0839,   0.5491,   1.9846],
    8.         [ -5.0909,  10.1805,  -8.2528,  -9.2240,  -1.8044,  -4.0296,  -8.2297,
    9.           -3.1828,  -5.9361,  -4.8410]], grad_fn=)

    6，保存模型

    最优模型默认保存在 trainer.checkpoint_callback.best_model_path 的目录下，可以直接加载。

    1. print(trainer.checkpoint_callback.best_model_path)
    2. print(trainer.checkpoint_callback.best_model_score)
    1. lightning_logs/version_10/checkpoints/epoch=8-step=15470.ckpt
    2. tensor(0.0376, device='cuda:0')
    1. model_clone = Model.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)
    2. trainer_clone = pl.Trainer(max_epochs=3,gpus=1) 
    3. result = trainer_clone.test(model_clone,data_module.test_dataloader())
    4. print(result)
    1. --------------------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9887999892234802, 'test_loss': 0.03627564385533333}
    4. --------------------------------------------------------------------------------
    5. [{'test_acc': 0.9887999892234802, 'test_loss': 0.03627564385533333}]

    三，训练加速技巧

    下面重点介绍pytorch_lightning 模型训练加速的一些技巧。

    • 1，使用多进程读取数据(num_workers=4)

    • 2，使用锁业内存(pin_memory=True)

    • 3，使用加速器(gpus=4,strategy="ddp_find_unused_parameters_false")

    • 4，使用梯度累加(accumulate_grad_batches=6)

    • 5，使用半精度(precision=16,batch_size=2*batch_size)

    • 6，自动搜索最大batch_size(auto_scale_batch_size='binsearch')

    （注：过大的batch_size对模型学习是有害的。）

    详细原理，可以参考：

    https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html

    我们将训练代码封装成如下脚本形式，方便后面测试使用。

    1. %%writefile mnist_cnn.py
    2. import torch 
    3. from torch import nn 
    4. from argparse import ArgumentParser
    5.  
    6. import torchvision 
    7. from torchvision import transforms as T
    8. from torchvision.datasets import MNIST
    9. from torch.utils.data import DataLoader,random_split
    10. import pytorch_lightning as pl
    11. from torchmetrics import Accuracy
    12.  
    13. #================================================================================
    14. # 一，准备数据
    15. #================================================================================
    16.  
    17. class MNISTDataModule(pl.LightningDataModule):
    18.     def __init__(self, data_dir: str = "./minist/", 
    19.                  batch_size: int = 32,
    20.                  num_workers: int =4,
    21.                  pin_memory:bool =True):
    22.         super().__init__()
    23.         self.data_dir = data_dir
    24.         self.batch_size = batch_size
    25.         self.num_workers = num_workers
    26.         self.pin_memory = pin_memory
    27.  
    28.     def setup(self, stage = None):
    29.         transform = T.Compose([T.ToTensor()])
    30.         self.ds_test = MNIST(self.data_dir, download=True,train=False,transform=transform)
    31.         self.ds_predict = MNIST(self.data_dir, download=True, train=False,transform=transform)
    32.         ds_full = MNIST(self.data_dir, download=True, train=True,transform=transform)
    33.         self.ds_train, self.ds_val = random_split(ds_full, [55000, 5000])
    34.  
    35.     def train_dataloader(self):
    36.         return DataLoader(self.ds_train, batch_size=self.batch_size,
    37.                           shuffle=True, num_workers=self.num_workers,
    38.                           pin_memory=self.pin_memory)
    39.  
    40.     def val_dataloader(self):
    41.         return DataLoader(self.ds_val, batch_size=self.batch_size,
    42.                           shuffle=False, num_workers=self.num_workers,
    43.                           pin_memory=self.pin_memory)
    44.  
    45.     def test_dataloader(self):
    46.         return DataLoader(self.ds_test, batch_size=self.batch_size,
    47.                           shuffle=False, num_workers=self.num_workers,
    48.                           pin_memory=self.pin_memory)
    49.  
    50.     def predict_dataloader(self):
    51.         return DataLoader(self.ds_predict, batch_size=self.batch_size,
    52.                           shuffle=False, num_workers=self.num_workers,
    53.                           pin_memory=self.pin_memory)
    54.     
    55.     @staticmethod
    56.     def add_dataset_args(parent_parser):
    57.         parser = ArgumentParser(parents=[parent_parser], add_help=False)
    58.         parser.add_argument('--batch_size', type=int, default=32)
    59.         parser.add_argument('--num_workers', type=int, default=4)
    60.         parser.add_argument('--pin_memory', type=bool, default=True)
    61.         return parser
    62.  
    63. #================================================================================
    64. # 二，定义模型
    65. #================================================================================
    66.  
    67. net = nn.Sequential(
    68.     nn.Conv2d(in_channels=1,out_channels=32,kernel_size = 3),
    69.     nn.MaxPool2d(kernel_size = 2,stride = 2),
    70.     nn.Conv2d(in_channels=32,out_channels=64,kernel_size = 5),
    71.     nn.MaxPool2d(kernel_size = 2,stride = 2),
    72.     nn.Dropout2d(p = 0.1),
    73.     nn.AdaptiveMaxPool2d((1,1)),
    74.     nn.Flatten(),
    75.     nn.Linear(64,32),
    76.     nn.ReLU(),
    77.     nn.Linear(32,10)
    78. )
    79.  
    80. class Model(pl.LightningModule):
    81.     
    82.     def __init__(self,net,learning_rate=1e-3):
    83.         super().__init__()
    84.         self.save_hyperparameters()
    85.         self.net = net
    86.         self.train_acc = Accuracy()
    87.         self.val_acc = Accuracy()
    88.         self.test_acc = Accuracy() 
    89.         
    90.         
    91.     def forward(self,x):
    92.         x = self.net(x)
    93.         return x
    94.     
    95.     
    96.     #定义loss
    97.     def training_step(self, batch, batch_idx):
    98.         x, y = batch
    99.         preds = self(x)
    100.         loss = nn.CrossEntropyLoss()(preds,y)
    101.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    102.     
    103.     #定义各种metrics
    104.     def training_step_end(self,outputs):
    105.         train_acc = self.train_acc(outputs['preds'], outputs['y']).item()    
    106.         self.log("train_acc",train_acc,prog_bar=True)
    107.         return {"loss":outputs["loss"].mean()}
    108.     
    109.     #定义optimizer,以及可选的lr_scheduler
    110.     def configure_optimizers(self):
    111.         return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
    112.     
    113.     def validation_step(self, batch, batch_idx):
    114.         x, y = batch
    115.         preds = self(x)
    116.         loss = nn.CrossEntropyLoss()(preds,y)
    117.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    118.  
    119.     def validation_step_end(self,outputs):
    120.         val_acc = self.val_acc(outputs['preds'], outputs['y']).item()    
    121.         self.log("val_loss",outputs["loss"].mean(),on_epoch=True,on_step=False)
    122.         self.log("val_acc",val_acc,prog_bar=True,on_epoch=True,on_step=False)
    123.     
    124.     def test_step(self, batch, batch_idx):
    125.         x, y = batch
    126.         preds = self(x)
    127.         loss = nn.CrossEntropyLoss()(preds,y)
    128.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    129.     
    130.     def test_step_end(self,outputs):
    131.         test_acc = self.test_acc(outputs['preds'], outputs['y']).item()    
    132.         self.log("test_acc",test_acc,on_epoch=True,on_step=False)
    133.         self.log("test_loss",outputs["loss"].mean(),on_epoch=True,on_step=False)
    134.     
    135.     @staticmethod
    136.     def add_model_args(parent_parser):
    137.         parser = ArgumentParser(parents=[parent_parser], add_help=False)
    138.         parser.add_argument('--learning_rate', type=float, default=1e-3)
    139.         return parser
    140.     
    141.  
    142. #================================================================================
    143. # 三，训练模型
    144. #================================================================================
    145.     
    146. def main(hparams):
    147.     pl.seed_everything(1234)
    148.     
    149.     data_mnist = MNISTDataModule(batch_size=hparams.batch_size,
    150.                                  num_workers=hparams.num_workers)
    151.     
    152.     model = Model(net,learning_rate=hparams.learning_rate)
    153.     
    154.     ckpt_callback = pl.callbacks.ModelCheckpoint(
    155.         monitor='val_loss',
    156.         save_top_k=1,
    157.         mode='min'
    158.     )
    159.     early_stopping = pl.callbacks.EarlyStopping(monitor = 'val_loss',
    160.                    patience=3,
    161.                    mode = 'min')
    162.     
    163.     trainer = pl.Trainer.from_argparse_args( 
    164.         hparams,
    165.         max_epochs=10,
    166.         
    167.         callbacks = [ckpt_callback,early_stopping]
    168.     ) 
    169.  
    170.     
    171.     
    172.     if hparams.auto_scale_batch_size is not None:
    173.         #搜索不发生OOM的最大batch_size
    174.         max_batch_size = trainer.tuner.scale_batch_size(model,data_mnist,
    175.                         mode=hparams.auto_scale_batch_size)
    176.         data_mnist.batch_size = max_batch_size
    177.         
    178.         #等价于
    179.         #trainer.tune(model,data_mnist)
    180.         
    181.     
    182.     #gpus=0, #单CPU模式
    183.     #gpus=1, #单GPU模式
    184.     #num_processes=4,strategy="ddp_find_unused_parameters_false", #多CPU(进程)模式
    185.     #gpus=4,strategy="dp", #多GPU(dp速度提升效果一般)
    186.     #gpus=4,strategy=“ddp_find_unused_parameters_false" #多GPU(ddp速度提升效果好)
    188.     trainer.fit(model,data_mnist)
    189.     result = trainer.test(model,data_mnist,ckpt_path='best')
    191. if __name__ == "__main__":
    192.     parser = ArgumentParser()
    193.     parser = MNISTDataModule.add_dataset_args(parser)
    194.     parser = Model.add_model_args(parser)
    195.     parser = pl.Trainer.add_argparse_args(parser)
    196.     hparams = parser.parse_args()
    197.     main(hparams)

    1，使用多进程读取数据(num_workers=4)

    使用多进程读取数据，可以避免数据加载过程成为性能瓶颈。

    • 单进程读取数据(num_workers=0, gpus=1): 1min 18s

    • 多进程读取数据(num_workers=4, gpus=1): 59.7s

    1. %%time
    2. #单进程读取数据(num_workers=0)
    3. !python3 mnist_cnn.py --num_workers=0 --gpus=1
    1. ------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9857000112533569, 'test_loss': 0.04885349050164223}
    4. --------------------------------------------------------------------------------
    5.  
    6. CPU times: user 4.67 s, sys: 2.14 s, total: 6.81 s
    7. Wall time: 2min 50s
    1. %%time
    2. #多进程读取数据(num_workers=4)
    3. !python3 mnist_cnn.py --num_workers=4 --gpus=1
    1. ---------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9764000177383423, 'test_loss': 0.0820135846734047}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|███████████████████████████████| 313/313 [00:01<00:00, 163.40it/s]
    6. CPU times: user 1.56 s, sys: 647 ms, total: 2.21 s
    7. Wall time: 59.7 s

    2，使用锁业内存(pin_memory=True)

    锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘）

    因此锁业内存比非锁业内存读写效率更高，copy到GPU上也更快速。

    当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。

    因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

    • 非锁业内存存储数据(pin_memory=False, gpus=1): 1min

    • 锁业内存存储数据(pin_memory=True, gpus=1): 59.5s

    1. %%time
    2. #非锁业内存存储数据(pin_memory=False)
    3. !python3 mnist_cnn.py --pin_memory=False --gpus=1
    1. ----------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9812999963760376, 'test_loss': 0.06231774762272835}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|███████████████████████████████| 313/313 [00:01<00:00, 171.69it/s]
    6. CPU times: user 1.59 s, sys: 619 ms, total: 2.21 s
    7. Wall time: 1min
    1. %%time
    2. #锁业内存存储数据(pin_memory=True)
    3. !python3 mnist_cnn.py --pin_memory=True --gpus=1
    1. ---------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9757999777793884, 'test_loss': 0.08017424494028091}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|███████████████████████████████| 313/313 [00:01<00:00, 174.58it/s]
    6. CPU times: user 1.54 s, sys: 677 ms, total: 2.22 s
    7. Wall time: 59.5 s

    3，使用加速器(gpus=4,strategy="ddp_find_unused_parameters_false")

    pl 可以很方便地应用单CPU、多CPU、单GPU、多GPU乃至多TPU上训练模型。

    以下几种情况训练耗时统计如下：

    • 单CPU: 2min 17s

    • 单GPU:  59.4 s

    • 4个GPU(dp模式): 1min

    • 4个GPU(ddp模式): 38.9 s

    一般情况下，如果是单机多卡，建议使用 ddp模式，因为dp模式需要非常多的data和model传输，非常耗时。

    1. %%time
    2. #单CPU
    3. !python3 mnist_cnn.py --gpus=0
    1. -----------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9790999889373779, 'test_loss': 0.07223792374134064}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|████████████████████████████████| 313/313 [00:05<00:00, 55.95it/s]
    6. CPU times: user 2.67 s, sys: 740 ms, total: 3.41 s
    7. Wall time: 2min 17s
    1. %%time
    2. #单GPU
    3. !python3 mnist_cnn.py --gpus=1
    1. ---------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9778000116348267, 'test_loss': 0.06929327547550201}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|███████████████████████████████| 313/313 [00:01<00:00, 171.04it/s]
    6. CPU times: user 1.83 s, sys: 488 ms, total: 2.32 s
    7. Wall time: 1min 3s
    1. %%time
    2. #多GPU，dp模式（为公平比较，batch_size=32*4）
    3. !python3 mnist_cnn.py --gpus=4 --strategy="dp" --batch_size=128
    1. ------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9790999889373779, 'test_loss': 0.06855566054582596}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|██████████████████████████████████| 79/79 [00:02<00:00, 38.55it/s]
    6. CPU times: user 1.2 s, sys: 553 ms, total: 1.75 s
    7. Wall time: 1min
    1. %%time
    2. #多GPU，ddp模式
    3. !python3 mnist_cnn.py --gpus=4 --strategy="ddp_find_unused_parameters_false"
    1. ---------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9732000231742859, 'test_loss': 0.08606339246034622}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|██████████████████████████████████| 79/79 [00:00<00:00, 85.79it/s]
    6. CPU times: user 784 ms, sys: 387 ms, total: 1.17 s
    7. Wall time: 38.9 s

    4，使用梯度累加(accumulate_grad_batches=6)

    梯度累加就是累加多个batch的梯度，然后用累加的梯度更新一次参数，使用梯度累加相当于增大batch_size.

    由于更新参数的计算量略大于简单梯度求和的计算量（对于大部分优化器而言），使用梯度累加会让速度略有提升。

    • 4个GPU(ddp模式): 38.9 s

    • 4个GPU(ddp模式)+梯度累加: 36.9 s

    1. %%time
    2. #多GPU，ddp模式, 考虑梯度累加
    3. !python3 mnist_cnn.py --accumulate_grad_batches=6 --gpus=4 --strategy="ddp_find_unused_parameters_false"
    1. ----------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9603000283241272, 'test_loss': 0.1400066614151001}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|██████████████████████████████████| 79/79 [00:00<00:00, 89.10it/s]
    6. CPU times: user 749 ms, sys: 402 ms, total: 1.15 s
    7. Wall time: 36.9 s

    5，使用半精度(precision=16)

    通过precision可以设置 double (64), float (32), bfloat16 ("bf16"), half (16) 精度的训练。

    默认是float(32) 标准精度，bfloat16 ("bf16")是混合精度。

    如果选择 half(16) 半精度，并同时增大batch_size为原来2倍, 通常训练速度会提升3倍左右。

    1. %%time 
    2. #半精度
    3. !python3 mnist_cnn.py --precision=16 --batch_size=64 --gpus=1

    6，自动搜索最大batch_size(auto_scale_batch_size="power")

    !python3 mnist_cnn.py --auto_scale_batch_size="power"  --gpus=1

    四，训练涨分技巧

    pytorch_lightning 可以非常容易地支持以下训练涨分技巧:

    • SWA(随机参数平均): 调用pl.callbacks.stochastic_weight_avg.StochasticWeightAveraging实现。

    • CyclicLR(学习率周期性调度策略): 设置 lr_scheduler 为 torch.optim.lr_scheduler.CyclicLR实现。

    • auto_lr_find最优学习率发现: 设置 pl.Trainer(auto_lr_find = True)实现。

    参考论文:

    • Cyclical Learning Rates for Training Neural Networks 【https://arxiv.org/pdf/1506.01186.pdf】

    • Averaging Weights Leads to Wider Optima and Better Generalization【https://arxiv.org/abs/1803.05407】

    我们将代码整理成如下形式，以便后续测试使用。

    1. %%writefile mnist_cnn.py
    2. import torch 
    3. from torch import nn 
    4. from argparse import ArgumentParser
    5. import numpy as np 
    6.  
    7. import torchvision 
    8. from torchvision import transforms as T
    9. from torchvision.datasets import MNIST
    10. from torch.utils.data import DataLoader,random_split
    11. import pytorch_lightning as pl
    12. from torchmetrics import Accuracy
    13.  
    14. #================================================================================
    15. # 一，准备数据
    16. #================================================================================
    17.  
    18. class MNISTDataModule(pl.LightningDataModule):
    19.     def __init__(self, data_dir: str = "./minist/", 
    20.                  batch_size: int = 32,
    21.                  num_workers: int =4,
    22.                  pin_memory:bool =True):
    23.         super().__init__()
    24.         self.data_dir = data_dir
    25.         self.batch_size = batch_size
    26.         self.num_workers = num_workers
    27.         self.pin_memory = pin_memory
    28.  
    29.     def setup(self, stage = None):
    30.         transform = T.Compose([T.ToTensor()])
    31.         self.ds_test = MNIST(self.data_dir, download=True,train=False,transform=transform)
    32.         self.ds_predict = MNIST(self.data_dir, download=True, train=False,transform=transform)
    33.         ds_full = MNIST(self.data_dir, download=True, train=True,transform=transform)
    34.         ds_train, self.ds_val = random_split(ds_full, [59000, 1000])
    35.         #为加速训练，随机取10000个
    36.         indices = np.arange(59000)
    37.         np.random.shuffle(indices)
    38.         self.ds_train = torch.utils.data.dataset.Subset(
    39.             ds_train,indices = indices[:3000]) 
    40.  
    41.     def train_dataloader(self):
    42.         return DataLoader(self.ds_train, batch_size=self.batch_size,
    43.                           shuffle=True, num_workers=self.num_workers,
    44.                           pin_memory=self.pin_memory)
    45.  
    46.     def val_dataloader(self):
    47.         return DataLoader(self.ds_val, batch_size=self.batch_size,
    48.                           shuffle=False, num_workers=self.num_workers,
    49.                           pin_memory=self.pin_memory)
    50.  
    51.     def test_dataloader(self):
    52.         return DataLoader(self.ds_test, batch_size=self.batch_size,
    53.                           shuffle=False, num_workers=self.num_workers,
    54.                           pin_memory=self.pin_memory)
    55.  
    56.     def predict_dataloader(self):
    57.         return DataLoader(self.ds_predict, batch_size=self.batch_size,
    58.                           shuffle=False, num_workers=self.num_workers,
    59.                           pin_memory=self.pin_memory)
    60.     
    61.     @staticmethod
    62.     def add_dataset_args(parent_parser):
    63.         parser = ArgumentParser(parents=[parent_parser], add_help=False)
    64.         parser.add_argument('--batch_size', type=int, default=32)
    65.         parser.add_argument('--num_workers', type=int, default=8)
    66.         parser.add_argument('--pin_memory', type=bool, default=True)
    67.         return parser
    68.  
    69. #================================================================================
    70. # 二，定义模型
    71. #================================================================================
    72.  
    73. net = nn.Sequential(
    74.     nn.Conv2d(in_channels=1,out_channels=32,kernel_size = 3),
    75.     nn.MaxPool2d(kernel_size = 2,stride = 2),
    76.     nn.Conv2d(in_channels=32,out_channels=64,kernel_size = 5),
    77.     nn.MaxPool2d(kernel_size = 2,stride = 2),
    78.     nn.Dropout2d(p = 0.1),
    79.     nn.AdaptiveMaxPool2d((1,1)),
    80.     nn.Flatten(),
    81.     nn.Linear(64,32),
    82.     nn.ReLU(),
    83.     nn.Linear(32,10)
    84. )
    85.  
    86. class Model(pl.LightningModule):
    87.     
    88.     def __init__(self,net,
    89.                  learning_rate=1e-3,
    90.                  use_CyclicLR = False,
    91.                  epoch_size=500):
    92.         super().__init__()
    93.         self.save_hyperparameters() #自动创建self.hparams
    94.         self.net = net
    95.         self.train_acc = Accuracy()
    96.         self.val_acc = Accuracy()
    97.         self.test_acc = Accuracy() 
    98.         
    99.         
    100.     def forward(self,x):
    101.         x = self.net(x)
    102.         return x
    103.     
    104.     
    105.     #定义loss
    106.     def training_step(self, batch, batch_idx):
    107.         x, y = batch
    108.         preds = self(x)
    109.         loss = nn.CrossEntropyLoss()(preds,y)
    110.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    111.     
    112.     #定义各种metrics
    113.     def training_step_end(self,outputs):
    114.         train_acc = self.train_acc(outputs['preds'], outputs['y']).item()    
    115.         self.log("train_acc",train_acc,prog_bar=True)
    116.         return {"loss":outputs["loss"].mean()}
    117.     
    118.     #定义optimizer,以及可选的lr_scheduler
    119.     def configure_optimizers(self):
    120.         optimizer = torch.optim.RMSprop(self.parameters(), lr=self.hparams.learning_rate)
    121.         if not self.hparams.use_CyclicLR:
    122.             return optimizer 
    123.  
    124.         max_lr = self.hparams.learning_rate
    125.         base_lr = max_lr/4.0
    126.         scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer,
    127.             base_lr=base_lr,max_lr=max_lr,
    128.             step_size_up=5*self.hparams.epoch_size,cycle_momentum=False)
    129.         self.print("set lr = "+str(max_lr))
    130.         
    131.         return ([optimizer],[scheduler])
    132.     
    133.     def validation_step(self, batch, batch_idx):
    134.         x, y = batch
    135.         preds = self(x)
    136.         loss = nn.CrossEntropyLoss()(preds,y)
    137.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    138.  
    139.     def validation_step_end(self,outputs):
    140.         val_acc = self.val_acc(outputs['preds'], outputs['y']).item()    
    141.         self.log("val_loss",outputs["loss"].mean(),on_epoch=True,on_step=False)
    142.         self.log("val_acc",val_acc,prog_bar=True,on_epoch=True,on_step=False)
    143.     
    144.     def test_step(self, batch, batch_idx):
    145.         x, y = batch
    146.         preds = self(x)
    147.         loss = nn.CrossEntropyLoss()(preds,y)
    148.         return {"loss":loss,"preds":preds.detach(),"y":y.detach()}
    149.     
    150.     def test_step_end(self,outputs):
    151.         test_acc = self.test_acc(outputs['preds'], outputs['y']).item()    
    152.         self.log("test_acc",test_acc,on_epoch=True,on_step=False)
    153.         self.log("test_loss",outputs["loss"].mean(),on_epoch=True,on_step=False)
    154.     
    155.     @staticmethod
    156.     def add_model_args(parent_parser):
    157.         parser = ArgumentParser(parents=[parent_parser], add_help=False)
    158.         parser.add_argument('--learning_rate', type=float, default=7e-3)
    159.         parser.add_argument('--use_CyclicLR', type=bool, default=False)
    160.         return parser
    161.     
    162.  
    163. #================================================================================
    164. # 三，训练模型
    165. #================================================================================
    166.     
    167. def main(hparams):
    168.     pl.seed_everything(1234)
    169.     
    170.     data_mnist = MNISTDataModule(batch_size=hparams.batch_size,
    171.                                  num_workers=hparams.num_workers)
    172.     data_mnist.setup()
    173.     epoch_size = len(data_mnist.ds_train)//data_mnist.batch_size
    174.     
    175.     model = Model(net,learning_rate=hparams.learning_rate,
    176.                   use_CyclicLR = hparams.use_CyclicLR,
    177.                   epoch_size=epoch_size)
    178.     
    179.     ckpt_callback = pl.callbacks.ModelCheckpoint(
    180.         monitor='val_acc',
    181.         save_top_k=3,
    182.         mode='max'
    183.     )
    184.     
    185.     early_stopping = pl.callbacks.EarlyStopping(monitor = 'val_acc',
    186.                    patience=16,
    187.                    mode = 'max')
    188.     callbacks = [ckpt_callback,early_stopping]
    189.     if hparams.use_swa:
    190.         callbacks.append(pl.callbacks.StochasticWeightAveraging())
    191.         
    192.     trainer = pl.Trainer.from_argparse_args( 
    193.         hparams,
    194.         max_epochs=1000,
    195.         callbacks = callbacks) 
    196.  
    197.     
    198.     print("hparams.auto_lr_find=",hparams.auto_lr_find)
    199.     if hparams.auto_lr_find:
    200.         
    201.         #搜索学习率范围
    202.         lr_finder = trainer.tuner.lr_find(model,
    203.           datamodule = data_mnist,
    204.           min_lr=1e-08,
    205.           max_lr=1,
    206.           num_training=100,
    207.           mode='exponential',
    208.           early_stop_threshold=4.0
    209.           )
    210.         lr_finder.plot() 
    211.         lr = lr_finder.suggestion()
    212.         model.hparams.learning_rate = lr 
    213.         print("suggest lr=",lr)
    214.         
    215.         del model 
    216.         
    217.         hparams.learning_rate = lr
    218.         model = Model(net,learning_rate=hparams.learning_rate,
    219.                   use_CyclicLR = hparams.use_CyclicLR,
    220.                   epoch_size=epoch_size)
    221.         
    222.         #等价于
    223.         #trainer.tune(model,data_mnist)
    224.         
    225.  
    226.     trainer.fit(model,data_mnist)
    227.     train_result = trainer.test(model,data_mnist.train_dataloader(),ckpt_path='best')
    228.     val_result = trainer.test(model,data_mnist.val_dataloader(),ckpt_path='best')
    229.     test_result = trainer.test(model,data_mnist.test_dataloader(),ckpt_path='best')
    230.     
    231.     print("train_result:\n")
    232.     print(train_result)
    233.     print("val_result:\n")
    234.     print(val_result)
    235.     print("test_result:\n")
    236.     print(test_result)
    237.     
    238.     
    239.  
    240. if __name__ == "__main__":
    241.     parser = ArgumentParser()
    242.     parser.add_argument('--use_swa', default=False, type=bool)
    243.     parser = MNISTDataModule.add_dataset_args(parser)
    244.     parser = Model.add_model_args(parser)
    245.     parser = pl.Trainer.add_argparse_args(parser)
    246.     hparams = parser.parse_args()
    247.     main(hparams)

    1，SWA 随机权重平均 (pl.callbacks.stochastic_weight_avg.StochasticWeightAveraging)

    • 平凡方式训练：test_acc = 0.9581000208854675

    • SWA随机权重：test_acc = 0.963100016117096

    1. #平凡方式训练
    2. !python3 mnist_cnn.py --gpus=2 --strategy="ddp_find_unused_parameters_false"
    1. ------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9581000208854675, 'test_loss': 0.14859822392463684}
    4. --------------------------------------------------------------------------------
    1. #使用SWA随机权重
    2. !python3 mnist_cnn.py --gpus=2 --strategy="ddp_find_unused_parameters_false" --use_swa=True
    1. -----------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.963100016117096, 'test_loss': 0.18146753311157227}
    4. --------------------------------------------------------------------------------

    2，CyclicLR学习率调度策略（torch.optim.lr_scheduler.CyclicLR）

    • 平凡方式训练：test_acc = 0.9581000208854675

    • SWA随机权重：test_acc = 0.963100016117096

    • SWA随机权重 + CyClicLR学习率调度策略: test_acc = 0.9688000082969666

    !python3 mnist_cnn.py --gpus=2 --strategy="ddp_find_unused_parameters_false" --use_swa=True --use_CyclicLR=True
    1. ------------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9688000082969666, 'test_loss': 0.11470437049865723}
    4. --------------------------------------------------------------------------------

    3, 最优学习率搜索(auto_lr_find=True)

    • 平凡方式训练：test_acc = 0.9581000208854675

    • SWA随机权重：test_acc = 0.963100016117096

    • SWA随机权重 + CyClicLR学习率调度策略: test_acc = 0.9688000082969666

    • SWA随机权重 + CyClicLR学习率调度策略 + 最优学习率搜索：test_acc = 0.9693999886512756

    !python3 mnist_cnn.py --gpus=1  --auto_lr_find=True --use_swa=True --use_CyclicLR=True
    1. ---------------------------------------------------------------
    2. DATALOADER:0 TEST RESULTS
    3. {'test_acc': 0.9693999886512756, 'test_loss': 0.11024412512779236}
    4. --------------------------------------------------------------------------------
    5. Testing: 100%|███████████████████████████████| 313/313 [00:02<00:00, 137.85it/s]

    以上。

    万水千山总是情，点个在看行不行？😋

    公众号后台回复关键词：pl，获取本文jupyter notebook源代码。

    04fcaeb92eb3c446af1697f3d653408b.png

  • 相关阅读:
    数智未来 持续创新 | 易趋受邀出席CIAS 2022中国数智汽车峰会
    关于视频超分辨率的学习-day1
    Python输入输出、遍历文件夹（input、stdin、os.path）
    盒子模型详解
    数据结构六：线性表之顺序栈的设计
    第三届 “鹏城杯”（初赛）
    Java专题训练——21天学习挑战赛
    v73.结构
    企业自己申报高企会遇到哪些问题，如何处理？
    win10 环境下Python 3.8按装fastapi paddlepaddle 进行身份证及营业执照的识别2
  • 原文地址：https://blog.csdn.net/Python_Ai_Road/article/details/126446693