Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

用自己的数据,重新训练BigGAN!

安妮 夏乙 假装发自 凹非寺

量子位 出品 | 公众号 QbitAI

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

“不,你不想。”

每次被BigGAN史上最强”的效果吸引,想要用其他数据集训练一番,脑海深处都会响起这样一个声音。

就仿佛DeepMind团队训练BigGAN用的512个TPU,齐刷刷发出不怀好意的嘲笑。

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

现在,

好消息来了!

BigGAN一作Andrew Brock新鲜放出了只需4-8个GPU即可训练的新版模型,这一次,他用的是PyTorch

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

虽然“只需”令人嫉妒,但比起原来动辄128个、512个TPU的硬件需求,简直就扶贫济困,平民之光。

网友们一夜之间把Brock这条Twitter转了100多次,在Reddit论坛上也夸这个新实现”fantastic”、”awesome”、”cool”,表示一定要玩。

更让人期待的是,打破了128个TPU的门槛之后,用其他数据集训练的BigGAN一定会层出不穷啊!风景、waifu指日可待。

另外,吸引网友们的不仅仅是算力需求的降低,还有:

等你的PyTorch版实现等了好久了!

摆脱算力束缚

Brock开源的低成本版PyTorch实现,内容包括训练、测试、采样所需的脚本、把THHub原版BigGAN生成器权重迁移到PyTorch的脚本,还有完整的预训练检查点。

除了现有的ImageNet预训练模型之外,他还打算放出Places-365数据集上预训练的模型。

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

如果你想用自己的数据来训练BigGAN,可以选择从头开始训练,也可以在ImageNet预训练模型的基础上微调。无论如何,只需4-8个GPU。

Brock还在Reddit论坛上说,其实一个GPU也行,但就是要花的时间太长了。

原版是128-512个TPU,新版是4-8个GPU,这之间有巨大的算力差距,Brock到底是怎样用Pytorch完整复现BigGAN的?

大思路就是:靠梯度累加,在小硬件上使用大批量(batch)。Brock已经将整个复现过程详细记录了下来。

我们先来看看这个低成本版BigGAN怎么用

开始之前,需要准备三样训练工具:

  • PyTorch(版本1.0.1)
  • tqdm,numpy,scipy和h5py
  • ImageNet训练集

首先,你可以选择为你的目标数据集准备一个预处理HDF5版本,以实现更快的I/O,这需要通过修改和运行计算FID需要的Inception moment。运行这行代码就可以:

sh scripts/utils/prepare_data.sh

默认情况下,它会假设你的ImageNet训练集下载到根目录的data位置下,并将以128×128像素的分辨率准备缓存的HDF5。

在脚本文件夹中,有多个bash脚本可以用来训练不同批大小的BigGAN。这些代码默认你无法访问完整的TPU pod,因此,会运用梯度累加技术来假冒大批量。

具体来说,先在多个小批量上将梯度平均,然后在N次累计后,再执行优化程序步骤。

在默认情况下,launch_BigGAN_bs256x8.sh脚本训练一个全尺寸96通道的BigGAN模型,用批大小为256的8个梯度累加,得到总体大小为2048的批量。

用的硬件少了,难免要多花点时间。按上面的默认情况,在128×128像素的ImageNet上训练BigGAN,用8个Tesla V100进行全精度的无张量训练,大约需要15天完成150000次迭代。

正确的流程是,你首先需要做的就是设置可以支持的最大的批大小,确定后再修改脚本,使批大小×梯度累加的数量=所需的批大小总量(BigGAN默认为2048)。

需要注意,梯度累加意味着需要将预估的奇异值(singular value,SV)和批归一化(Batch Normalization,BN)统计量提高8倍。

这意味着,BN的统计值更接近于静止状态,相较之下奇异值的估计值往往更加准确。

因此,Brock在默认情况下用测试模式下的G来衡量标准,即用BN的运行状态预估值来替代去计算静态统计值。不过,最终repo也支持静态统计,虽然可能会导致前期的梯度累加失效,但问题也不大。

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

标准和示例

在训练期间,这个脚本会输出具有训练标准和测试标准的日志,并且将保存模型权重/优化的程序参数的多个副本,并且每次保存权重时,都将生成示例和插值。

日志文件夹中存储了处理这些日志的脚本,并将结果用Matlab绘制出来。

训练过程完成后,你可以用sample.py代码块生成其他示例和插值,还能改变截断值、批大小、静态统计累计值等进行测试。

默认情况下,所有内容都保存到权重/示例/日志/数据文件夹中,repo被假定到与它们在同一个文件夹里了。你可以用—base_root代码块将这些参数指向不同的基础文件夹,或者用各自的参数选择每个基础文件夹的特定位置。

在Brock小哥给出的实现中,包含了BigGAN-deep的运行脚本,但没有用它完全训练成一个模型,是未经测试的一个脚本。

此外,Brock还给出了包括在CIFAR上运行的脚本,以及在ImageNet上运行SA-GAN和SN-GAN的脚本。SA-GAN是假设用4个TitanX训练时的脚本,是在批大小为128时加2个梯度累加的情况下进行的。

用自己的数据微调预训练模型

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

我们前边也提到过,除了从头开始训练,还可以拿Brock预训练的模型来微调。

要运行的脚本和从头训练一样,但是要加上—resume。

继续进行中断了的训练也是用这个脚本。

要准备微调用的数据集,需要添加到datasets.py中,并在utils.py里调整convenience dicts,例如dset_dict、imsize_dict、root_dict、nclass_dict、classes_per_sheet_dict等等设置,这样才能为数据集创建合适的元数据。

然后,重复prepar_data.sh的流程。

我有一块1080Ti,行吗?

低成本版一出,引来不少网友欢呼,很多一看就来自学神们。

他们很是关注Brock所用的方法,比如同步的批归一化(synchronized batchnorm)、梯度累加(gradient accumulation)技术等等,还打算用到自己的研究里。

当然,更多人关心的是:

用这个“官方非官方”PyTorch版实现训练,成本究竟怎么样?

按照Brock举的例子计算,8块Tesla V100训练15天,在亚马逊云服务AWS上的价格是8800多美元。

也就是说,如果你自己本地没有那么多GPU,还是挺贵……

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

于是,评论区除了欢呼的大神们之外,也冒出了很多贫穷的声音:

哇哦好赞啊!但是你说得好像4-8块GPU多便宜似的哈哈哈哈。

然而我只有一块1080ti……

在Colab里哭。(量子位注:就是蹭免费GPU的Colab啦)

也有人真诚提问:

我有一块2080ti,你估计半个星期能训练到收敛吗?还是需要几周?

Brock给出的答案不容乐观:一块2080ti啊,那得好几个月了。

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

也就是说,如果预算不充沛,还是拿预训练模型微调的好。

不过,评论区也有高人,为新实现找了个绝佳的使用场景:

那些之前买显卡挖矿的人,终于有比较酷的事情可干了。

这一提议顿时有人响应:

哈……我有12块英伟达1080,就是因为这个。

如果你也有挖矿遗留下来的GPU,那就一切好办了。

转投PyTorch体验怎么样?

除了省算力,“终于有了完整PyTorch版实现”也是网友们非常关心的。

毕竟,Brock作为一名(曾经的)DeepMind实习生,改用PyTorch实属难得。这一版本“官方非官方”实现是他和Alex Andonian一起完成的,而并没有提到BigGAN另外两位作者。

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

 BigGAN一作Andrew Brock

他使用PyTorch体验怎样?如何看待在PyTorch和TensorFlow中实现BigGAN的差异?

在Reddit上,也有网友向Brock好奇地提出了这个问题。

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

Brock坦言,用PyTorch训练BigGAN,是种非常可爱的体验

他认为PyTorch让很多事情都变得容易多了,比如组成层、测试、调试、实现模型并行性等操作,一切都可以信手拈来。

整个实现的过程中,当然也遇到了一些没想到的困难,不过Brock说都是他自己的错。比如直到开始检查梯度时,才发现错位的ReLU已经造成了大面积破坏。

此外,他遇到的最大挑战其实是算力不足:因为Brock只有多所大学的共享服务器可用,因此调试/测试周期也拉长了不少。每次训练只能运行24小时,等待下一次开始的时间,则不止24小时了。

传送门

Bye Bye TPU,4个GPU训练“史上最强”GAN!作者开源PyTorch模型

最后,附上GitHub地址:

https://github.com/ajbrock/BigGAN-PyTorch

还有论文:

Large Scale GAN Training for High Fidelity Natural Image Synthesis

https://arxiv.org/abs/1809.11096

版权所有,未经授权不得以任何形式转载及使用,违者必究。