|
“不，你不想。” 每次被BigGAN史上最强”的效果吸引，想要用其他数据集训练一番，脑海深处都会响起这样一个声音。 就仿佛DeepMind团队训练BigGAN用的512个TPU，齐刷刷发出不怀好意的嘲笑。
Brock开源的低成本版PyTorch实现，内容包括训练、测试、采样所需的脚本、把THHub原版BigGAN生成器权重迁移到PyTorch的脚本，还有完整的预训练检查点。 除了现有的ImageNet预训练模型之外，他还打算放出Places-365数据集上预训练的模型。
如果你想用自己的数据来训练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也支持静态统计，虽然可能会导致前期的梯度累加失效，但问题也不大。
标准和示例 在训练期间，这个脚本会输出具有训练标准和测试标准的日志，并且将保存模型权重/优化的程序参数的多个副本，并且每次保存权重时，都将生成示例和插值。 日志文件夹中存储了处理这些日志的脚本，并将结果用Matlab绘制出来。 训练过程完成后，你可以用sample.py代码块生成其他示例和插值，还能改变截断值、批大小、静态统计累计值等进行测试。 默认情况下，所有内容都保存到权重/示例/日志/数据文件夹中，repo被假定到与它们在同一个文件夹里了。你可以用—base_root代码块将这些参数指向不同的基础文件夹，或者用各自的参数选择每个基础文件夹的特定位置。 在Brock小哥给出的实现中，包含了BigGAN-deep的运行脚本，但没有用它完全训练成一个模型，是未经测试的一个脚本。 此外，Brock还给出了包括在CIFAR上运行的脚本，以及在ImageNet上运行SA-GAN和SN-GAN的脚本。SA-GAN是假设用4个TitanX训练时的脚本，是在批大小为128时加2个梯度累加的情况下进行的。 用自己的数据微调预训练模型
我们前边也提到过，除了从头开始训练，还可以拿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，还是挺贵……
还有论文：
| 
发表于 2019-3-23 19:21:11
