SAST Weekly 是由电子工程系学生科协推出的科技系列推送,内容涵盖信息领域技术科普、研究前沿热点介绍、科技新闻跟进探索等多个方面,帮助同学们增长姿势,开拓眼界,每周更新,欢迎关注!欢迎愿意分享知识的同学投稿至 [email protected] , 期待你的作品!
相信已经有不少同学接触过Pytorch了(媒认课上还见到了九字班的同学),作为目前最流行的深度学习框架之一,Pytorch本身已经具备了入门简单,代码简洁等优点,它可以让你把精力专注在在实现自己的模型上,很适合学生、研究人员来使用。
尽管Pytorch中需要用户自己实现的功能已经非常少了,但是除了定义自己的模型之外,我们还需要在代码中加入导入、加载数据集,手写训练、交叉验证循环,保存、绘制loss曲线等等写起来重复度高并且还会让你的代码变得不那么优美的内容(主要是写起来麻烦)。那么,有没有一种——
Pytorch Lightning可以帮助你把写Pytorch变得更加轻松愉快[小恐龙点赞]
官方文档是这样描述它的
Lightning is a way to organize your PyTorch code to decouple the science code from the engineering. It's more of a PyTorch style-guide than a framework.
简单来说,你只需要关注你的模型的实现,也就是Research Code
,其他的Engineering code
都交给Lightning的Trainer来实现
先用一张图来看看Pytorch Lightning能为你做什么
具体到代码,你可以省去图中蓝色的部分
而原本的定义Pytorch模型的代码可以直接用于LightningModule。事实上,你只需要用另一种风格来重新组织你的Pytorch代码而不需要学习新的概念。
接下来我们用一个MNIST手写数字识别的例子来体验下Lightning。
首先模型的父类要换成LightningModule
__init__
和forward
方法可以直接使用不需要修改。
与Pytorch不同的是,你需要把加载数据以及创建DataLoader的代码写在模型中,需要定义如下几个方法
事实上我们只是把加载数据集的代码放到了prepare_data中,把创建dataloader的代码分别放到了train_dataloader
和val_dataloader
中
接下来我们要把原本的training loop中最内层的代码写到training_step
和validation_step
中
training_step
和validation_step
中的返回值都是一个字典,其中training_step的返回值必须包含loss属性,log是你需要输出的数据的字典,而validation_step
的返回值字典保存的是你想要在每次validation结束输出的数据,如下
返回的字典中progress_bar也需要是一个字典,对应显示在进度条上的值(稍后展示),log同上作为输出的数据被保存下来
最后,我们还需要定义一个optimizer
这样一来一个LightningModule就完成了,而他的训练异常简单,你只需要
不需要手写training loop,也不需要自己加入进度条,一切都由Trainer帮你完成
可以看到进度条右侧显示出了我们设置的validation的结果
如果你想使用GPU训练,也不必像原来一样一个个设置,只需要
至于我们在step中设置要输出的log,它们被按照tensorboard(可替换其他可视化工具)的格式保存在了./lightning_logs
中,你可以直接使用tensorboard来查看我们保存的log的曲线
如果你在模型中定义了test_dataloader
, test_step
和test_epoch_end
,你就可以直接进行test
以上简单地展示了Pytorch Lightning的使用,已经可以明显地看出他在pytorch上做的改进:你的代码条理更加清晰,你不需要自己完成诸如training loop的工程性强研究意义低的代码,你也不需要为可视化花心思。
除了改变了写代码的风格之外,其余没有与pytorch不同的地方,实际上你的代码中使用的依然是pytorch里的方法和函数,只是把他们交给了Lightning来组织、运行、输出
Lightning对RNN、GAN等也有着很好的支持,Trainer可设置的参数也十分丰富,在此不再作介绍,有兴趣的同学可以参考官方文档进行深入了解
参考资料:
官方文档:
https://pytorch-lightning.readthedocs.io/en/latest/
From PyTorch to PyTorch Lightning — A gentle introduction:
https://towardsdatascience.com/from-pytorch-to-pytorch-lightning-a-gentle-introduction-b371b7caaf09
撰稿:李煜泽
审核:孙志尧