Pytorch-lightning+Default Checkpoint+Hydra+Tensorboard实践

在深度学习实践中，如何构建一个系统往往需要阅读大量的开源项目、自己总结经验。这里根据我的习惯，总结了常用框架和工具链的使用逻辑。也可搜索Havard CS197系统学习。

pytorch-lightning

基本逻辑被LightningModule（简称LM）封装net，里边forward方法负责inference，在此基础上定义training_step 利用forward的输出logits计算loss，以及valid_step等等hook去控制。剩下一切项目配置都在Trainer中指定。

Trainer中可以指定参数，还可以引入callbacks和plugin。前者例如ckpt，后者例如amp和ddp strategy。logger是单独的一个参数，也可以继承后传入。

1. 关于log：

普通的tensorboard SummaryWriter可以add_text 添加文本信息，还可以传入数值信息化成曲线。但是PL封装的没有add_text接口，记录文字很不方便，所以可以全局维护一个loguru的日志去单独记录文字（或者封装两个logger 成为一个大的logger）这也是为什么深度学习系统一般是两个logger。

使用方法：

        1. pl_loggers.TensorBoardLogger() 作为参数传进Trainer，成为LM成员

        2. LM中使用self.log / self.log_dict 即可访问到它，向日志写入内容（key-alue形式，Value

来自于train step输出的东西）图片等需要不同的方法

        3. log记录的东西也能被其他组件捕捉到，例如ModelCheckpoint的filename，可以自定义格式化参数，你log的loss可以作为文件名

        4. metrics：self.log能够智慧的累加取平均。如果想定义更复杂的算子，使用torch.Metrics，它也能自动累加、平均，而且比self.log还支持了ddp自动reduce

        5. 其他：

                如何用epoch而非step做x-axis链接，重写log ‘step’变量即可

                想看lr使用 callbacks = LearnRateMonitor，它添加所有的optimizer的lr曲线

        6. Trying to infer the `batch_size` from an ambiguous collection. The batch size we found is To avoid any miscalculations, use self.log(...,batch_size=batch_size)`[issue] 。这个warning说的是log可以自动累加计算metric，但是batch_size不明确的话“平均”操作会意义不统一，要求在log时显式的声明batch size

2. 关于checkpoint（文档）

内置的ModelCheckpoint是包办lr, opt, stateful datamodule和 stateful callbacks的（什么是stateful）。根据官方文档，想继续训练只需创建一个model，其他都不用创建。不过这听起来不大科学，据我实验，至少在1.8 之前的版本，datamodule和callbacks都需要手动传入，不过它们的state可以保存，比如“ModelCheckpoint的保存目录”这个state，继续训练后还是会保持之前的，而不是新建文件夹。

model = LitModel()
trainer = Trainer()

# automatically restores model, epoch, step, LR schedulers, apex, etc...
trainer.fit(model, ckpt_path="some/path/to/my_checkpoint.ckpt")

这是一个ckpt的例子，其中callbacks就保存了ModelCheckpoint这些callback。

dict_keys(['epoch', 'global_step', 'pytorch-lightning_version', 'state_dict', 
'loops', 'callbacks', 'optimizer_states', 'lr_schedulers', 'hparams_name', 
'hyper_parameters'])

想只load模型inference：

model = MyLightningModule.load_from_checkpoint

save_hyperparameters是LM init时候传入的参数。即便你用了hydra，传入的直接是实例化之后的torch module，lightning也可以这个module写到ckpt的'hyper_parameters'字段。建议把配置参数或者实例化过的一切都存入ckpt，省的每次去修改config文件了，直接跑就行.

但是，如果参数被写进save_hyperparameters的ignore列表，就不会记录进ckpt，那么load_from_checkpoint这个方法就不会读入它，于是就会报错缺少初始化参数。解决办法是还用torch的老办法，写入state_dict。（其中torch.load的是字典，load_state_dict是把字典加载到模型权重里）

checkpoint = torch.load(checkpoint_path)

checkpoint.keys()
>>> dict_keys(['epoch', 'global_step', 'pytorch-lightning_version', 'state_dict', 'callbacks', 'optimizer_states', 'lr_schedulers', 'native_amp_scaling_state', 'hparams_name', 'hyper_parameters', 'hparams_type'])
#传统的torch存储的pth的内容
>>> dict_keys(['meta', 'state_dict', 'lr_scheduler', 'optimizer', 'amp'])

cfg = OmegaConf.to_object(checkpoint['hyper_parameters'])
cfg = DictConfig(cfg)
model = build_model(cfg)
model.load_state_dict(checkpoint['state_dict'])

综上，比较好的实践是：LM只传入cfg，完全保存 hyperparameters，构建backbone和loss、metrics等操作放在LM内部。

Strict=False

load_state_dict 和 load_from_checkpoint都有strict参数

with open(cfg.get("ckpt_path"),'rb') as f:
     st = torch.load(f)['state_dict']
     model.load_state_dict(state_dict=st,strict=False)
trainer.fit(model=model, datamodule=datamodule)

 只有当PL读取了整个ckpt路径时，才会恢复optim，lr，state，和所有超参数。只load state dict不会再读取超参数了。

超参数重载

参数的应用场景有两个：要么构建模型时传入，要么发给trainner，前者作为hyperparmeters保存，用load_from_checkpoint重载，后者用Trainer()里边的字典重载

eg，如果你训了10个epoch，又想再来5个，但是之前参数传入的是10，这样一开始训练就会stop了。那怎么重载呢？可以只设定 Trainer(max_epoches=15)。参考issue：

https://github.com/Lightning-AI/lightning/issues/2823

部分参数重载

如果我只想重载weights，lr，optmizer等等使用新的重新训练可以吗？有没有可选择的重载方式？

更新：可以使用load_state_dict。pl能包办optimizer的前提是trainer.fit或者trainer初始化时，传入了ckptpath的参数。但load_state_dict只告知了模型的状态，并没有告知optimizer和lr scheduler的状态。

---------------------------------------------------------------------------------

2023.1.8日为止，没有。github上该issue还是open状态，等待新的feature引入吧！现在的办法有2种：

        1. 在load_from_checkpoint之前修改ckpt。如果想改存ckpt的文件夹，可以修改ModelCheckpoint的状态。但是把lr和opti的state dict删除掉可能会导致error。这种等同于只load weights作为初始化参数了。

        2. issue的第三个回答，使用configure_optimizer函数，在一定epoch开始之后reset opti和lrscheduler

https://github.com/Lightning-AI/lightning/issues/5339

存储规则：

ModelCheckpoint这个；类在初始化时候有个‘monitor’可以监控self.log下来的参数。比如你记录的accu，可设它是accu，就会根据这个保存topk。如果只是想傻瓜的每轮都保存，使用参数：

save_top_k = -1
every_n_train_steps = cfg.CHECKPOINT_INTERVAL,

继续训练：

搞定了上边所有，就可以继续训练了。主要关注的有4个：模型结构和权重；有状态的callback；hydra；logger的路径。前两个已经被lightning搞定了，后两个需要自己重新设定路径。logger的root path通常会依赖于一个由${hydra:runtime.output_dir}决定的变量。而hydra可以由config/hydra自由配置

3. 关于optimizer：

train step：可以renturn 一个scalar，也可以是带‘loss’的key的字典

optimizer：在LM 里用configure** 定义，

manual optimize可以控制如强化学习等高度自由的优化过程

配套的lr scheduler也可以在里边调整，返回一个字典

return [optimizer], [{'scheduler': scheduler, 'interval': 'epoch'}] # interval step or epoch

容易犯的bug：如果interval写成了 ‘step’，是每个batch都会更新lr，lr快速逼近0！

4. 关于AMP：文档

 Pytorch半精度网络训练需要注意哪些问题？ | w3c笔记

有2个backend，PL推荐native后端。amp_level是专门给apex后端设的。1.10建议用plugin而不是参数配置了。注意有的操作必须在32bit下完成！

warn

Do not cast anything to other dtypes manually using torch.autocast or tensor.half() when using native precision because this can bring instability.

也就是说不用把任何东西导成half。当然float和double的差异还是要控制一下的，在dataset和loss计算中通常手动处理

【BUG】现在amp会有过度调整lr的bug，此bug源于torch内部，和PL无关！

        这个bug会产生warning：

UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate

它是在开始几个epoch调用lr_scheduler.step() when the scaler skips optimizer.step()带来的。可以看下列issue

Mixed precision: scheduler and optimizer are called in the wrong order · Issue #5558 · Lightning-AI/lightning · GitHub

Fix LR scheduler getting called excessively with amp by akihironitta · Pull Request #11542 · Lightning-AI/lightning · GitHub

5. 关于config：

hydra层次化读取，omegaconf解析引用（yaml里边的value可以引用其他的key，hydra读出来都是str，需要这个包解析（Omegaconf.resolve））

注意，yaml中defaults list是从上往下重载的，有同名的取最后一个的值

Installation — OmegaConf 2.1.3.dev0 documentation

Tutorial: Learning Hydra for configuring ML experiments - Simone Scardapane

比较好的实践是：

1. 把不太容易变动的东西写进default yaml
2. 把正式训练时的所有需要调参的写进shell命令行里进行重载
3. 做过拟合实验或者debug所需要重载的所有参数放进一个experiment模块，在主yaml下边进行重载  experiment: exp_overfit

我之前觉得hydra文档写的烂，可能错怪它了，很多用法都是omegaconf的设定，要去读他们的文档！ 

一些用例：

• 在命令行中override一个submodule，而不是一个attribute: 不要用"."而要用"\" ，例如 model/head/lane=lane_segment_model
• 在命令行中override一个list的第7个submodule的一个属性(transform是那个list)：data.transform.7.attr=xxxx
• yaml中可以定义 ：attr: &in_dims 7, 那个in_dims就是变量，可以在同一个文件中用*引用
• yaml可以插入omegaconf支持的表达式，例如${eval:${a}*${b}}, omegaconf参考文档

6. 关于wandb、tensorboard和可视化：

6.1 tensorboard：它的曲线经过smooth，所以有时候要看value，而不是smooth，这才是真实的log值，后者是tb插值得到的！

在服务器上docker内部需要加  --bind_all

6.2 wandb：

7. 关于DDP：

lightning 在validation阶段是不会自动聚合不同gpu上的loss的，所以你需要手动all gather一下，具体做法在这里。不过，torchmetrics包（用法在这里）支持了自动all gather，此时不用手动操作了。顺便，lr是和optimizer绑定的，lr调度时候的“monitor”依据的指标就是这里同步下来的才行？

syncbatchnorm：  多gpu时如果单卡的batchsize太小，统计的方差也会不准，所以把多卡算出来的方差放一起。torch.nn.SyncBatchNorm.convert_sync_batchnorm 可以把所有的模块统一转换，不过PL中只需Trainer的参数设一下就好了。这里主要是训练时的动作，推理时不需要。

其他trick

2. torch.cuda.synchronize()：应对异步情况，比如 tensor.to(non_blocking=True)送数据，用它同步

3. 类似的深度好文   pytorch_lightning 全程笔记 - 知乎

4. 可以通过 len(tensor.shape) 判断