江苏龙网

  1. 首页 > 资讯 > 科技 > >

PyTorch|让PyTorch更轻便，这款框架值得拥有！GitHub 6.6k星


白交 发自 凹非寺
量子位 报道 | 公众号 QbitAI
一直以来 ， PyTorch就以简单又好用的特点 ， 广受AI研究者的喜爱 。
但是 ， 一旦任务复杂化 ， 就可能会发生一系列错误 ， 花费的时间更长 。
于是 ， 就诞生了这样一个“友好”的PyTorch Lightning 。
PyTorch|让PyTorch更轻便，这款框架值得拥有！GitHub 6.6k星
本文插图

直接在GitHub上斩获6.6k星 。
PyTorch|让PyTorch更轻便，这款框架值得拥有！GitHub 6.6k星
本文插图

首先 ， 它把研究代码与工程代码相分离 ， 还将PyTorch代码结构化 ， 更加直观的展现数据操作过程 。
这样 ， 更加易于理解 ， 不易出错 ， 本来很冗长的代码一下子就变得轻便了 ， 对AI研究者十分的友好 。
话不多说 ， 我们就来看看这个轻量版的“PyTorch” 。
关于Lightning
Lightning将DL/ML代码分为三种类型：研究代码、工程代码、非必要代码 。
针对不同的代码 ， Lightning有不同的处理方式 。
这里的研究代码指的是特定系统及其训练方式 ， 比如GAN、VAE ， 这类的代码将由LightningModule直接抽象出来 。
我们以MNIST生成为例 。
l1 = nn.Linear(...)
l2 = nn.Linear(...)
decoder = Decoder()
x1 = l1(x)
x2 = l2(x2)
out = decoder(features, x)
loss = perceptual_loss(x1, x2, x) + CE(out, x)
而工程代码是与培训此系统相关的所有代码 ， 比如提前停止、通过GPU分配、16位精度等 。
我们知道 ， 这些代码在大多数项目中都相同 ， 所以在这里 ， 直接由Trainer抽象出来 。
model.cuda(0)
x = x.cuda(0)
distributed = DistributedParallel(model)
with gpu_zero:
download_data()
dist.barrier()
剩下的就是非必要代码 ， 有助于研究项目 ， 但是与研究项目无关 ， 可能是检查梯度、记录到张量板 。 此代码由Callbacks抽象出来 。
# log samples
z = Q.rsample()
generated = decoder(z)
self.experiment.log('images', generated)
此外 ， 它还有一些的附加功能 ， 比如你可以在CPU ， GPU ， 多个GPU或TPU上训练模型 ， 而无需更改PyTorch代码的一行；你可以进行16位精度训练 ， 可以使用Tensorboard的五种方式进行记录 。
这样说 ， 可能不太明显 ， 我们就来直观的比较一下PyTorch与PyTorch Lightning之间的差别吧 。
PyTorch与PyTorch Lightning比较
直接上图 。
PyTorch|让PyTorch更轻便，这款框架值得拥有！GitHub 6.6k星
本文插图

我们就以构建一个简单的MNIST分类器为例 ， 从模型、数据、损失函数、优化这四个关键部分入手 。
模型
首先是构建模型 ， 本次设计一个3层全连接神经网络 ， 以28×28的图像作为输入 ， 将其转换为数字0-9的10类的概率分布 。
PyTorch|让PyTorch更轻便，这款框架值得拥有！GitHub 6.6k星
本文插图

两者的代码完全相同 。 意味着 ， 若是要将PyTorch模型转换为PyTorch Lightning ， 我们只需将nn.Module替换为pl.LightningModule
也许这时候 ， 你还看不出这个Lightning的神奇之处 。 不着急 ， 我们接着看 。
数据
接下来是数据的准备部分 ， 代码也是完全相同的 ， 只不过Lightning做了这样的处理 。
它将PyTorch代码组织成了4个函数 ， prepare_data、train_dataloader、val_dataloader、test_dataloader
PyTorch|让PyTorch更轻便，这款框架值得拥有！GitHub 6.6k星


推荐阅读


上一篇：5G|苹果iPhone 12系列电池缩水？5G冲浪一天一充可能都有些吃紧

下一篇：超能网|高考完了该装机了：AORUS推出毕业相关优惠活动，更有Z390以旧换新