一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

新闻 前端
对于PyTorch开发者来说,Google Colab上的TPU资源不能用,恐怕是最遗憾的事情了。Google的云计算资源眼看着不能用,转TensorFlow又是不可能的。

 本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。

对于PyTorch开发者来说,Google Colab上的TPU资源不能用,恐怕是最遗憾的事情了。Google的云计算资源眼看着不能用,转TensorFlow又是不可能的。

过去一直有PyTorch用户试图在Colab上薅羊毛,但是都没有太成功的。

现在福利来了,一个叫做Pytorch Lightning的项目,可以让你几乎修改代码的情况下用上TPU。

[[316583]]

Pytorch Lightning已经上传到PyPI,因此只需一行代码就能安装这个软件。

  1. pip install pytorch-lightning 

该项目的开发者William Falcon说,PyTorch Lightning是他在纽约大学和FAIR做博士生时研发,专门为从事AI研究的专业研究人员和博士生创建的。

几乎无需修改代码

首先让我们来看一个MNIST图像分类网络的搭建,PyTorch的原始代码和修改后的PyTorch Lightning代码几乎无异。

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

我们只需将nn.Module替换为pl.LightningModule即可。

作者表示,相比切换框架,用这种方法重构原来的代码只需数小时的时间。

PyTorch Lightning所做的是将科学代码与工程代码分离,只需将研究代码重构为LightningModule格式(科学),Lightning将自动完成其余部分(工程)。

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

使用方法

PyTorch Lightning具体该如何使用,作者Falcon还是以MNIST图像分类网络为例,介绍从收集数据到训练再到验证、测试的全过程。

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

准备数据集阶段分为下载图片、转换、分割数据集、打包四个步骤。二者代码大致相同,只是将PyTorch代码组织为4个函数:

prepare_data:此函数负责处理下载数据,确保使用多个GPU时,不会下载多个数据集或对数据进行双重操作。

train_dataloader,val_dataloader,test_dataloader:每一个都负责返回相应数据集的数据拆分。以这种方式进行构造训练、验证、测试集,可以让你非常清楚如何操作数据。

接下来是优化器的选择,比如选择Adam。两者的代码还是几乎完全相同,不过后者把pytorch_model改成了self。

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

至于损失函数,对于n向分类,要使用交叉熵损失。二者代码又几乎一致相同,后者多出一个self。

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

在训练上,PyTorch Lightning的代码更简洁一点。在PyTorch中,你需要自己编写for循环,这意味着你必须记住要以正确的顺序调用正确的东西,可能会导致错误。

而PyTorch Lightning将样板格式抽象化,但方框中的内容保持不变,增加了代码的可读性和可重复性。

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

后面验证推理的部分不再赘述。

总之,PyTorch Lightning有这些优点:

代码结构化;与PyTorch源代码几乎完全相同;随着项目复杂性的提升,代码的大部分内容无需修改;保留了PyTorch的灵活性。

新增高级功能,连Pytorch本尊都没有

除了以上的一些特性外,PyTorch Lightning还加入了许多高级功能,让你体验到PyTorch本身不具备的一些优点。

比如更清晰直观的训练进度条:

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

用TensorBoard日志记录代码运行全过程:

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

PyTorch Lightning还支持TensorBoard之外的其它5种工具记录日志:

一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

与TensorBoard,MLFlow完全集成,并支持任何日志记录模块。

甚至还有一个内置的分析工具,告诉你训练过程中的瓶颈:

  1. trainer = Trainer(…, profiler=True) 
一行代码安装,TPU也能运行PyTorch,修改少量代码即可快速移植

PyTorch Lightning还有更多的可扩展性,在这里无法一一介绍,如果你正想要在TPU上运行自己的PyTorch代码,可以前去学习更详细的用法。

传送门

项目地址:

https://github.com/PyTorchLightning/pytorch-lightning

Colab演示:

https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3#scrollTo=dEeUzX_5aLrX

责任编辑:张燕妮 来源: 量子位
相关推荐

2020-04-01 08:51:47

注释字符代码

2020-07-15 09:40:37

代码Python浏览记录

2022-03-21 17:54:52

Linux操作系统浏览器

2016-12-02 08:53:18

Python一行代码

2014-02-12 13:43:50

代码并行任务

2017-04-05 11:10:23

Javascript代码前端

2022-04-09 09:11:33

Python

2023-09-12 10:10:57

开发者工具开源

2020-08-05 11:53:41

数据代码自动化

2021-08-31 09:49:37

CPU执行语言

2017-04-13 19:20:18

Python代码并行任务

2020-08-19 10:30:25

代码Python多线程

2020-09-09 16:00:22

Linux进程

2021-11-02 16:25:41

Python代码技巧

2022-04-11 11:38:44

Python代码游戏

2019-12-25 14:08:50

Pandas数据计算

2020-08-12 14:54:00

Python代码开发

2020-09-28 12:34:38

Python代码开发

2018-03-08 11:43:18

PandasTB级数据Spark

2022-02-08 22:18:10

Chrome插件服务器
点赞
收藏

51CTO技术栈公众号