一行代码安装，TPU也可以运行PyTorch，只需修改少量代码即可快速移植。对于PyTorch的开发者来说，谷歌Colab上的TPU资源无法使用大概是最遗憾的事情了。谷歌的云计算资源用不上，转用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-way分类，使用交叉熵损失。两者的代码几乎相同，后者多了一个自己。在训练方面，PyTorchLightning的代码更简洁一些。在PyTorch中，你需要自己编写for循环，这意味着你必须记住以正确的顺序调用正确的东西，这可能会导致错误。虽然PyTorchLightning抽象了样板格式，但框中的内容保持不变，增加了代码的可读性和可重复性。验证推理部分后面不再赘述。简而言之，PyTorchLightning具有以下优点：代码结构化；它几乎与PyTorch源代码相同；随着项目复杂度的增加，大部分代码不需要修改；保留了PyTorch的灵活性。新的高级功能，连Pytorch本身都没有除了上面的一些功能，PyTorchLightning还增加了很多高级功能，让你体验到一些PyTorch本身没有的优势。比如更清晰直观的训练进度条：使用TensorBoardlog记录运行代码的全过程：PyTorchLightning除了TensorBoard还支持其他五种工具记录日志：它与TensorBoard和MLFlow完全集成，支持任意记录模块。甚至还有一个内置的分析工具可以告诉你训练过程中的瓶颈：trainer=Trainer(...,profiler=True)如果你想运行自己的PyTorch代码，PyTorchLightning比这里可以介绍的具有更多的可扩展性关于TPU，大家可以去了解更详细的用法。传送门项目地址：https://github.com/PyTorchLightning/pytorch-lightningColabdemo：https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3#scrollTo=dEeUzX_5aLrX
