datasheet

QbitAI量子位

文章数:1711 被阅读:20240101

账号入驻

在PyTorch上用"Keras",分布式训练开箱即用,告别没完没了的Debug

2019-08-07
    阅读数:
鱼羊 发自 凹非寺 
量子位 报道 | 公众号 QbitAI

在开始一个新的机器学习项目时,难免要重新编写训练循环,加载模型,分布式训练……然后在Debug的深渊里看着时间哗哗流逝,而自己离项目核心还有十万八千里。

虽然这世上已经有了神器Keras,能用几条语句就轻松组建一个神经网络,但一想到它是站在Tensorflow的肩膀上,就让人不禁想起江湖中的那句传说:

PyTorch 真香!

那么为什么不做一个PyTorch上的Keras呢?

来自Facebook的Willian Falcon小哥决定一试,他搞了个包装器,把PyTorch中的各种通用配置全部包装在一起。

这个PyTorch轻量级包装器,就是PyTorch Lightning

有了这样一个快速研究框架,使用者只需关注核心训练和验证逻辑,繁琐的工程细节通通自动化一键完成,既能保证核心训练逻辑的正确性,又能保证最佳的实践体验。

像闪电一样迅疾

所以,Lightning到底有多好用?

在这张图中,灰色部分代表Lightning能自动完成的部分,而蓝色的部分则能够根据使用者的需求,被定义成任意底层模型,可以是你自己的新模型,也可以是预训练模型,fast.ai架构等等。

举几个例子好了。

比如说,梯度下降

原来,你需要这样:

# clear last step
optimizer.zero_grad()

# 16 accumulated gradient steps
scaled_loss = 0
for accumulated_step_i in range(16):
    out = model.forward()
    loss = some_loss(out,y)
    loss.backward()
    scaled_loss += loss.item()

# update weights after 8 steps. effective batch = 8*16
optimizer.step()

# loss is now scaled up by the number of accumulated batches
actual_loss = scaled_loss / 16

在Lightning里,这一整段代码不需要你自己敲了,只需输入以下两行代码:

trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

不仅如此,在Lightning里,想用上单个GPU,直接调用即可:

trainer = Trainer(gpus = [0])
trainer.fit(model)

使用能将内存占用减少一半的黑科技16位精度,不费吹灰之力:

trainer = Trainer(amp_level = ’02’, use_amp = False)
trainer.fit(model)

使用多个GPU进行分布式训练,so easy:

trainer = Trainer(gpus = [0123])
trainer.fit(model)



甚至是在1024个节点上以1024个GPU进行训练,也是开箱即用:

trainer = Trainer(nb_gpu_nodes=128, gpus=[01234567])

自动化功能还远不止这一些,以下模块,均包含其中:

想要训练闪电那么快的神经网络吗?Lightning简直为此量身定做。

此外,Lightning还和Tensorboard集成在了一起,可以轻松实现可视化学习。

只需记录实验路径:

from test_tube import Experiment
from pytorch-lightning import  Trainer

exp = Experiment(save_dir='/some/path')
trainer = Trainer(experiment=exp)
...

然后在该路径运行tensorboard即可:

tensorboard —logdir /some/path

食用方法

想要使用Lightning,需要完成两件事。

1、定义Lightning Model

这一步会花费掉比较长的时间。

Lightning Model是nn.Module的严格超类,它提供了与模型进行交互的标准界面。

启用Lightning Model最简单的方法是根据下面这个最小示例(minimal example)进行局部修改:

import os
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import torchvision.transforms as transforms

import pytorch_lightning as ptl

class CoolModel(ptl.LightningModule):

    def __init__(self):
        super(CoolModel, self).__init__()
        # not the best model...
        self.l1 = torch.nn.Linear(28 * 2810)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def my_loss(self, y_hat, y):
        return F.cross_entropy(y_hat, y)

    def training_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
        return {'loss': self.my_loss(y_hat, y)}

    def validation_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
        return {'val_loss': self.my_loss(y_hat, y)}

    def validation_end(self, outputs):
        avg_loss = torch.stack([x['val_loss'for x in outputs]).mean()
        return {'avg_val_loss': avg_loss}

    def configure_optimizers(self):
        return [torch.optim.Adam(self.parameters(), lr=0.02)]

    @ptl.data_loader
    def tng_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

    @ptl.data_loader
    def val_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

    @ptl.data_loader
    def test_dataloader(self):
        return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

2、拟合训练器

训练器能处理Lightning自动化部分的代码核心逻辑,它会在训练过程中提取出最佳实践。

基本的用法是像这样:

from pytorch_lightning import Trainermoder = LightningTemplate()trainer = Trainer()
trainer.fit(model)

只要确保它的正确执行,只需一个Trainer,计算集群(SLURM),Debug,分布式训练就通通不在话下了。

from pytorch_lightning import Trainer
from test_tube import Experiment

model = CoolModel()
exp = Experiment(save_dir=os.getcwd())

# train on cpu using only 10% of the data (for demo purposes)
trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)

# train on 4 gpus
# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3])

# train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)

# train (1 epoch only here for demo)
trainer.fit(model)

# view tensorflow logs 
print(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')
print('and going to http://localhost:6006 on your browser')


One More Thing

你可能会问,为什么要搞一个Lightning呢,用fast.ai不好吗?

作者小哥表示,Lightning和fast.ai之间就没什么好比的,fast.ai面向有志于进入深度学习领域的新手,而Lightning面向的是ML领域中活跃的研究人员们。

就算真的要比,Lightning可是开箱即用的,不仅如此,在高性能计算、调试工具和可用性方面,小哥都对Lightning充满信心。他自信地甩出了三张对比表格:



嗯,PyTorch真香!

传送门

GitHub地址:
https://github.com/williamFalcon/pytorch-lightning

教学博客:
https://towardsdatascience.com/supercharge-your-ai-research-with-pytorch-lightning-337948a99eec


加入社群 | 与优秀的人交流

小程序 | 全类别AI学习教程


量子位 QbitAI · 头条号签约作者

վ'ᴗ' ի 追踪AI技术和产品新动态

喜欢就点「好看」吧 !




最新有关QbitAI的文章

About Us 关于我们 客户服务 联系方式 器件索引 网站地图 最新更新 手机版

站点相关: TI培训

北京市海淀区知春路23号集成电路设计园量子银座1305 电话:(010)82350740 邮编:100191

电子工程世界版权所有 京ICP证060456号 京ICP备10001474号 电信业务审批[2006]字第258号函 京公海网安备110108001534 Copyright © 2005-2018 EEWORLD.com.cn, Inc. All rights reserved