Pytorch Lightning
pl是一个轻量级的PyTorch库,用于对深度学习模型进行轻量化,但它并不提供样板,而仅仅是轻量化
它能对日常使用的样板代码进行抽象和自动化。
从上面我们可以发现 pl 的三个优势
- 通过抽象出样板工程代码,可以更容易地识别和理解ML代码。
- Lightning的统一结构使得在现有项目的基础上进行构建和理解变得非常容易。
- Lightning 自动化的代码是用经过全面测试、定期维护并遵循ML最佳实践的高质量代码构建的。
pl.LightningModule
pl.LightningModule
是Pytorch中的模型部分(nn.Module
)和训练逻辑(训练、验证、测试)部分的结合。
继承自该类的函数通常包含三个部分:
- 模型相关部分:
__init__
forward
- 优化器相关部分:
configure_optimizers
- 模型训练逻辑部分:
training_step
validation_step
test_step
模型相关部分一般涉及模型的构建,包括超参数定义、模型初始化、模型运行逻辑(forward
函数)
优化器相关部分一般涉及模型的优化器初始化,学习率的schedule设置等
训练逻辑部分一般是每个训练、验证、预测步骤需要进行的操作。
1 | import pytorch_lightning as pl |
pl 流程
PL的流程很简单,生产流水线,有一个固定的顺序:
初始化 def __init__(self)
-->训练training_step(self, batch, batch_idx)
--> 校验validation_step(self, batch, batch_idx)
--> 测试 test_step(self, batch, batch_idx)
. 就完事了,总统是实现这三个函数的重写。
当然,除了这三个主要的,还有一些其他的函数,为了方便我们实现其他的一些功能,因此更为完整的流程是在training_step 、validation_step、test_step 后面都紧跟着其相应的 training_step_end(self,batch_parts)
和training_epoch_end(self, training_step_outputs)
函数,当然,对于校验和测试,都有相应的*_step_end
和*_epoch_end
函数。因为校验和测试的*_step_end
函数是一样的,因此这里只以训练为例。
*_step_end
– 即每一个 * 步完成后调用
*_epoch_end
– 即每一个 * 的epoch 完成之后会自动调用
Train
训练主要是重写def training_setp(self, batch, batch_idx)
函数,并返回要反向传播的loss
即可,其中batch 即为从 train_dataloader 采样的一个batch的数据,batch_idx即为目前batch的索引。
1 | def training_setp(self, batch, batch_idx): |
Validation
设置校验的频率
每训练n个epochs 校验一次
默认为每1个epoch
校验一次,即自动调用validation_step()
函数
1 | trainer = Trainer(check_val_every_n_epoch=1) |
单个epoch内校验频率
当一个epoch 比较大时,就需要在单个epoch 内进行多次校验,这时就需要对校验的调动频率进行修改, 传入val_check_interval
的参数为float
型时表示百分比,为int
时表示batch
:
1 | # 每训练单个epoch的 25% 调用校验函数一次,注意:要传入float型数 |
校验和训练是一样的,重写def validation_step(self, batch, batch_idx)
函数,不需要返回值:
1 | def validation_step(self, batch, batch_idx): |
Test
在pytoch_lightning
框架中,test 在训练过程中是不调用的,也就是说是不相关,在训练过程中只进行training和validation,因此如果需要在训练过中保存validation的一些信息,就要放到validation中。
关于测试,测试是在训练完成之后的,因此这里假设已经训练完成:
1 | # 获取恢复了权重和超参数等的模型 |
数据集
数据集有两种实现方法:
当然,首先要自己先实现Dataset的定义,可以用现有的,例如MNIST等数据集,若用自己的数据集,则需要自己去继承torch.utils.data.dataset.Dataset
,自定义类,这一部分不再细讲,查其他的资料。
直接实现
直接实现是指在Model中重写def train_dataloader(self)
等函数来返回dataloader:
1 | class ExampleModel(pl.LightningModule): |
这样就完成了数据集和dataloader的编程了,注意,要先自己完成dataset的编写,或者用现有的normal数据集
自定义DataModule
这种方法是继承pl.LightningDataModule
来提供训练、校验、测试的数据。
1 | class MyDataModule(pl.LightningDataModule): |
使用
1 | dm = MyDataModule(args) |
模型保存
自动保存
Lightning 会自动保存最近训练的epoch的模型到当前的工作空间(or.getcwd()
),也可以在定义Trainer的时候指定:
1 | trainer = Trainer(default_root_dir='/your/path/to/save/checkpoints') |
当然,也可以关闭自动保存模型:
1 | trainer = Trainer(checkpoint_callback=False) |
ModelCheckpoint(call backs)
- 计算需要监控的量,例如校验误差:
loss
- 使用
log()
函数标记该要监控的量 - 初始化
ModelCheckpoint
回调,并设置要监控的量,下面有详细的描述 - 将其传回到
Trainer
中
1 | from pytorch_lightning.callbacks import ModelCheckpoint |
其中ModuleCheckpoint
的参数如下:
1 | CLASS pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint(filepath=None, monitor=None, verbose=False, save_last=None, save_top_k=None, save_weights_only=False, mode='auto', period=1, prefix='', dirpath=None, filename=None) |
参数说明:所有参数均为optional。
filepath
– 不建议使用,在后续版本中会被删除;保存的模型文件的路径,后面的参数会有另外两个参数来代替这个。
monitor
– 需要监控的量,string
类型。例如'val_loss'
(在training_step()
orvalidation_step()
函数中通过self.log(‘val_loss’, loss)进行标记);默认为None
,只保存最后一个epoch的模型参数,(我的理解是只保留最后一个epoch的模型参数,但是还是每训练完一个epoch之后会保存一次,然后覆盖上一次的模型)
verbose
:冗余模式,默认为False.
save_last
:bool
类型; 默认None
,当为True
时,表示在每个epoch 结果的时候,总是会保存一个模型last.ckpt
,也就意味着会覆盖保存,只会有一个文件保留。
save_top_k
:int
类型;当save_top_k==k
,根据monitor
监控的量,保存k
个最好的模型,而最好的模型是当monitor
监控的量最大时表示最好,还是最小时表示最好,在后面的参数mode
中进行设置。当save_top_k==0
时,不保存;当save_top_k==-1
时,保存所有的模型,即每个次保存模型不进行覆盖保存,全都保存下来;当save_top_k>=2
,并且在单个epoch内多次调用保存模型的函数,则模型的名字最后会追加版本号,从v0
开始。
mode
:string
类型,只能取{‘auto’, ‘min’, ‘max’}中的一个;当save_top_k!=0
时,保存模型时就会覆盖保存,如果monitor
监控的是val_loss
等越小就表示模型越好的,这个参数应该被设置成'min'
,当monitor
监控的是val_acc
(校验准确度)等越大就表示模型训练的越好的量,则应该设置成'max'
。auto
会自动根据monitor
的名字来判断(auto
模式是个人理解,可能会出错,例如你编程的时候,你就喜欢用val_loss
表示模型准确度这样就会导致保存的模型是最差的模型了)。
save_weights_only
:bool
类型;True
只保存模型权重(model.save_weights(flepath)
),否则保存整个模型。建议保存权重就可以了,保存整个模型会消耗更多时间和存储空间。
period
:int
类型。保存模型的间隔,单位为epoch,即隔多少个epoch自动保存一次。
prefix
:string
类型;保存模型文件的前缀。
dirpath
:string
类型。例如:dirpath='my/path_to_save_model/'
filename
:string
类型;前面就说过不建议使用filepath
变量,推荐使用dirpath+filename
的形式来作为模型路径。例如:
文件名会以epoch、val_loss、和其他的一些指标作为名称来保存
模型名称: my/path/epoch=2-val_loss=0.02-other_metric=0.03.ckpt
checkpoint_callback = ModelCheckpoint( ... , dirpath='my/path', ... filename='{epoch}-{val_loss:.2f}-{other_metric:.2f}' ... )
使用案例
1 | from pytorch_lightning import Trainer |
获取最好的模型
因为根据上面保存的参数,可能保存了多个模型,根据best_model_path
恢复最好的模型。
1 | from pytorch_lightning import Trainer |
手动保存模型
1 | from collections import deque |
1 | # 保存最新的路径 |
自定义文件名
format_checkpoint_name
(epoch, step, metrics, ver=None)
在上面的filename
参数中,定义了模型文件的保存格式,这个函数就是给其中的变量赋值的,返回string
类型,文件名
1 | >>> tmpdir = os.path.dirname(__file__) |
手动保存
除了自己保存,还可以手动保存和加载模型
1 | model = MyLightningModule(hparams) |
加载检查点
加载模型的权重、偏置和超参数:
1 | model = MyLightingModule.load_from_checkpoint(PATH) |
如果需要修改超参数,在写Module的时候进行覆盖:
1 | class LitModel(LightningModule): |
这样的话,可以如下恢复模型:
1 | # 例如训练的时候初始化in_dim=32, out_dim=10 |
load_from_checkpoint 方法
1 | LightningModule.load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, strict=True, **kwargs) |
恢复模型和Trainer
如果不仅仅是想恢复模型,而且还要接着训练,则可以恢复Trainer
1 | model = LitModel() |
训练辅助
Early Stopping
监控validation_step()
中某一个量,如果其不能再变得更优,则提前停止训练
1 | pytorch_lightning.callbacks.early_stopping.EarlyStopping(monitor='early_stop_on', min_delta=0.0, patience=3, verbose=False, mode='auto', strict=True) |
monitor (
str
) – 监控的量;默认为:early_stop_on
;可以通过self.log('var_name', val_loss)
来标记要监控的量
min_delta (float
) – 最小的改变量;默认:0.0;即当监控的量的绝对值变量量小于该值,则认为没有新的提升
patience (int
) - 默认:3;如果监控的量持续patience 个epoch没有得到更好的提升,则停止训练;
verbose (bool
) – 默认:False;
mode (str
) – {auto, min, max}中的一个,跟前面的ModelCheckpoint
中的mode
是一样的含义。如果monitor
监控的是val_loss
等越小就表示模型越好的,这个参数应该被设置成'min'
,当monitor
监控的是val_acc
(校验准确度)等越大就表示模型训练的越好的量,则应该设置成'max'
。auto
会自动根据monitor
的名字来判断(auto
模式是个人理解,可能会出错,例如你编程的时候,你就喜欢用val_loss
表示模型准确度这样就会导致保存的模型是最差的模型了)。
strict (bool
) – 默认True;如果监控器没有在validation_step()
函数中找到你监控的量,则强制报错,中止训练;
Logging
这里只涉及Tensorboard
, 其它有需要的可参考官方文档Logging,tensorboard 有两种基本的方法:一种是只适用于scaler,可直接使用self.log()
,另一种是图像、权重等。
1 | # 在定义Trainer对象的时候,传入tensorboardlogger |
注意如果是用anaconda的话,要先激活你的env,另外要注意的是,--logdir=my_log_dir/
, 这里的logdir要到version_0/
目录,该目录下保存有各种你log的变量的文件夹
1 | # 查看的方法跟tensorboard是一样的,在终端下 |
当然也可以继承LightningLoggerBase
类来自定义Logger,这个自己看官方文档
optimizer 和 lr_scheduler
当然,在训练过程中,对学习率的掌控也是非常重要的,合理设置学习率有利于提高效果,学习率衰减可查看四种学习率衰减方法。那在pytorch_lightning 中如何设置呢?其实跟pytorch是一样的,基本上不需要修改:
1 | # 重写configure_optimizers()函数即可 |
这样就OK了,只要在training_step()
函数中返回了loss
,就会自动反向传播,并自动调用loss.backward()
和optimizer.step()
和stepLR.step()
了
多优化器用于多模型等网络结构
当我们训练的是复杂的网络结构时,可能有多个模型,需要不同的训练顺序,不同的训练学习率等,这时候就需要设计多个优化器,并手动调用梯度反传函数
1 | # multiple optimizer case (e.g.: GAN) |
然后要关掉自动优化,这样就可以跟pytorch一样手动控制优化器的权重更新了,达到了跟pytorch一样可以进行复杂地更新顺序等地控制,同时pytorch lightning的优势还在,例如多GPU下batchnorm的参数同步等。
1 | # 在new Trainer对象的时候,把自动优化关掉 |
这时候 training_step()
函数也就不是直接返回loss 或者 字典了,而是不需要返回loss了,因为在该函数里就手动完成权重更新函数地调用, 另外需要注意的是:1. 不再使用loss.backward()
函数,改用self.manual_backward(loss, opt)
,就可以实现半精度训练。 2. 忽略optimizer_idx
参数。
1 | def training_step(self, batch, batch_idx, opt_idx): |
其他
多GPU训练
如果是CPU训练,在定义Trainer时不管gpus
这个参数就可以了,或者设置该参数为0
:
1 | trainer = pl.Trainer(gpus=0) |
而多GPU训练,也是很方便,只要将该参数设置为你要用的gpu数就可以,例如用4张GPU:
1 | trainer = pl.Trainer(gpus=4) |
而如果你有很多张GPU,但是要跟同学分别使用,只要在程序最前面设置哪些GPU可用就可以了,例如服务器有4张卡,但是你只能用0和2号卡:
1 | import os |
半精度训练
半精度训练也是Apex的一大特色,可以在几乎不影响效果的情况下降低GPU显存的使用率(大概50%),提高训练速度,现在pytorch_lightning 统统都给你,可以只要设置一下参数就可以:
1 | trainer = pl.Trainer(precision=16) |
累积梯度
默认情况是每个batch 之后都更新一次梯度,当然也可以N
个batch后再更新,这样就有了大batch size 更新的效果了,例如当你内存很小,训练的batch size 设置的很小,这时候就可以采用累积梯度:
1 | # 默认情况下不开启累积梯度 |
自动缩放batch_size
这方法还有很多限制,直接trainer.fit(model)
是无效的,感觉挺麻烦,不建议用
大的batch_size 通过可以获得更好的梯度估计。但同时也要更长的时间,另外,如果内存满了,电脑会卡住动不了。 'power'
– 从batch size 为1 开始翻倍地往上找,例如``1–>2 --> 4 --> …一直到内存溢出(out-of-memory, OOM);
binsearch`也是翻倍地找,直OOM,但是之后还要继续进行一个二叉搜索,找到一个更好的batch size。另外,搜索的batch size 最大不会超过数据集的尺寸。
1 | # 默认不开启 |
保存所有超参数到模型中
将所有的模型超参数都保存到模型中,恢复模型时再也不用自己去拖动恢复模型中的超参数了,这点是太有特色了:
1 | # 例如你传入的超参数字典为params_dict |
当然,对于无们训练的不同的模型,我们还是需要查看其超参数,可以通过将超参数字典保存到本地txt的方法,来以便后期查看
1 | def save_dict_as_txt(list_dict, save_dir): |
梯度剪裁
当需要避免发生梯度爆炸时,可以采用梯度剪裁的方法,这个梯度范数是通过所有的模型权重计算出来的:
1 | # 默认不剪裁 |
设置训练的最小和最大epochs
默认最小训练1个epoch,最大训练1000个epoch。
1 | trainer = Trainer(min_epochs=1, max_epochs=1000) |
小数据集
当我们的数据集过大或者当我们进行debug
时,不想要加载整个数据集,则可以只加载其中的一小部分:
默认是全部加载,即下面的参数值都为1.0
1 | # 参训练集、校验集和测试集分别只加载 10%, 20%, 30%,或者使用int 型表示batch |
其中比较需要注意的是训练集和测试集比例的设置,因为pytorch_lightning 每次validation和test时,都是会计算一个epoch,而不是一个step,因此在训练过程中,如果你的validation dataset比较大,那就会消耗大量的时间在validation上,而我们实际上只是想要知道在训练过程中,模型训练的怎么样了,不需要跑完整个epoch,因此就可以将limit_val_batches
设置的小一些。对于test,在训练完成后,如果我们不希望对所有的数据都进行test,也可以通过这个参数来设置。
另外,该框架有个参数num_sanity_val_steps
,用于设置在开始训练前先进行num_sanity_val_steps
个 batch的validation,以免你训练了一段时间,在校验的时候程序报错,导致浪费时间。该参数在获得trainer的时间传入:
1 | # 默认为2个batch的validation |
PL当前版本的BUG
该框架10月份才推出,有BUG才是正常的,下面给出我遇到的BUG(也有的是自己踩的坑):
多GPU CUDA设备不同步问题
在进行多GPU训练过程中,当完成一个epoch或者运行到epoch的指定百分比后,会进行validation过程,完成validation后,报了一个错:
1 | RuntimeError: All input tensors must be on the same device. Received cuda:2 and cuda:0 |
我在github上发起了一个issue,有人已经修复了这个bug,在后面新版的pytorch_lightning
中应该会被修改了,但是如果你用的时候也报这个错,可以通过下面的步骤进行修复:
DataLoader的问题
RuntimeError: DataLoader worker (pid(s) 6700, 10620) exited unexpectedly
这个问题一般是多GPU跑的时候才会出现,主要是加载DataLoader的时候,num_works=0
就可以了,另外,我在一个task里,设置的是num_works=8
是OK了,但是到了另一个task中,图像更大了,可能是内存不够,加载数据集特别特别特别慢,几乎不动。
如遇到这个报错,减小batch_size,设置num_works=0
,在定义trainer的时候,设置
1 | trainer = pl.Trainer(distributed_backend='ddp') |
一键导出模型训练过程中记录的loss
在训练过程当中,我们会用 self.log
函数去把训练的loss 和 校验的loss 等信息保存起来,下面就是我写的代码,一键导出某个训练模型中的所有log 数据,支持一键保存成excel,不同的loss 保存到同一个excel 文件中的不同表格,也提供了将excel 文件保存成 .mat 文件,方便用于matlab 的绘图放到论文中. 不需要 tensorflow, 有 tensorboard 就可以了