一行代码安装，TPU也能运行PyTorch，修改少量代码即可快速移植
晓查发自凹非寺量子位报道|公众号QbitAI
对于PyTorch开发者来说 ， GoogleColab上的TPU资源不能用 ， 恐怕是最遗憾的事情了 。 Google的云计算资源眼看着不能用 ， 转TensorFlow又是不可能的 。
文章图片
?
过去一直有PyTorch用户试图在Colab上薅羊毛 ， 但是都没有太成功的 。
现在福利来了 ， 一个叫做PytorchLightning的项目 ， 可以让你几乎修改代码的情况下用上TPU 。
【一行代码安装，TPU也能运行PyTorch，修改少量代码即可快速移植】
文章图片
PytorchLightning已经上传到PyPI ， 因此只需一行代码就能安装这个软件 。
pipinstallpytorch-lightning该项目的开发者WilliamFalcon说 ， PyTorchLightning是他在纽约大学和FAIR做博士生时研发 ， 专门为从事AI研究的专业研究人员和博士生创建的 。
几乎无需修改代码首先让我们来看一个MNIST图像分类网络的搭建 ， PyTorch的原始代码和修改后的PyTorchLightning代码几乎无异 。
文章图片
我们只需将nn.Module替换为pl.LightningModule即可 。
作者表示 ， 相比切换框架 ， 用这种方法重构原来的代码只需数小时的时间 。
PyTorchLightning所做的是将科学代码与工程代码分离 ， 只需将研究代码重构为LightningModule格式（科学） ， Lightning将自动完成其余部分（工程） 。
文章图片
使用方法PyTorchLightning具体该如何使用 ， 作者Falcon还是以MNIST图像分类网络为例 ， 介绍从收集数据到训练再到验证、测试的全过程 。
文章图片
准备数据集阶段分为下载图片、转换、分割数据集、打包四个步骤 。 二者代码大致相同 ， 只是将PyTorch代码组织为4个函数：
prepare_data：此函数负责处理下载数据 ， 确保使用多个GPU时 ， 不会下载多个数据集或对数据进行双重操作 。
train_dataloader ， val_dataloader ， test_dataloader：每一个都负责返回相应数据集的数据拆分 。 以这种方式进行构造训练、验证、测试集 ， 可以让你非常清楚如何操作数据 。
接下来是优化器的选择 ， 比如选择Adam 。 两者的代码还是几乎完全相同 ， 不过后者把pytorch_model改成了self 。
文章图片
至于损失函数 ， 对于n向分类 ， 要使用交叉熵损失 。 二者代码又几乎一致相同 ， 后者多出一个self 。
文章图片
在训练上 ， PyTorchLightning的代码更简洁一点 。 在PyTorch中 ， 你需要自己编写for循环 ， 这意味着你必须记住要以正确的顺序调用正确的东西 ， 可能会导致错误 。
而PyTorchLightning将样板格式抽象化 ， 但方框中的内容保持不变 ， 增加了代码的可读性和可重复性 。
文章图片
后面验证推理的部分不再赘述 。
总之 ， PyTorchLightning有这些优点：
代码结构化；与PyTorch源代码几乎完全相同；随着项目复杂性的提升 ， 代码的大部分内容无需修改；保留了PyTorch的灵活性 。
新增高级功能 ， 连Pytorch本尊都没有除了以上的一些特性外 ， PyTorchLightning还加入了许多高级功能 ， 让你体验到PyTorch本身不具备的一些优点 。
比如更清晰直观的训练进度条：
文章图片
用TensorBoard日志记录代码运行全过程：
文章图片
PyTorchLightning还支持TensorBoard之外的其它5种工具记录日志：
文章图片
与TensorBoard ， MLFlow完全集成 ， 并支持任何日志记录模块 。
甚至还有一个内置的分析工具 ， 告诉你训练过程中的瓶颈：
trainer=Trainer(…,profiler=True)
文章图片
PyTorchLightning还有更多的可扩展性 ， 在这里无法一一介绍 ， 如果你正想要在TPU上运行自己的PyTorch代码 ， 可以前去学习更详细的用法 。
传送门项目地址：
https://github.com/PyTorchLightning/pytorch-lightning
Colab演示：
https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3#scrollTo=dEeUzX_5aLrX
—完—
量子位QbitAI·***签约作者
关注我们 ， 第一时间获知前沿科技动态
推荐阅读
- 准备进场信号指标(副图)源码！
- 卫成印一行就重在点工程项目建设兴县调研
- 学院党委书记一行深入公共基础部指导返校复学工作
- 老司机经常用几个简单的驾驶代码词，赶紧学！
- 北海国家高新技术产业开发区领导一行莅临参观
- 潍坊市卫生健康委党组书记、主任吴海源一行到青岛西海岸新区考察交流
- 县长汪明君会见阳城企业集团有限公司执行董事赵志亮一行
- 治庸·铁军 | 盘龙公安分局5月21日全警动员组织开展“三查三告知”集中统一行动
- Win10推送Intel CPU专属补丁，无法取消更新，只能强制安装！
- 安装车顶行李箱的六大好处，老司机都知道