#量子位#TPU也能运行PyTorch，修改少量代码即可快速移植，一行代码安装晓查发自凹非寺量子位报道|公众号QbitA

晓查发自凹非寺量子位报道|公众号QbitAI
对于PyTorch开发者来说， GoogleColab上的TPU资源不能用，恐怕是最遗憾的事情了。 Google的云计算资源眼看着不能用，转TensorFlow又是不可能的。

文章图片
?
过去一直有PyTorch用户试图在Colab上薅羊毛，但是都没有太成功的。
现在福利来了，一个叫做PytorchLightning的项目，可以让你几乎修改代码的情况下用上TPU 。

文章图片
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本身不具备的一些优点。
比如更清晰直观的训练进度条：