tensorflow - 如何在 TF 2.0 / 1.14.0-eager 和自定义训练循环(梯度磁带)中执行梯度累积?
问题描述
背景:我有一个模型,我正在尝试移植它以TF 2.0
获得一些甜蜜的渴望执行,但我似乎无法弄清楚如何进行分布式训练(4 个 GPU)并同时执行梯度累积。
问题:
我需要能够使用带有梯度磁带的自定义训练循环,因为我有一个复杂的多模型问题(几个输入模型和输出模型一起训练),我不需要二阶梯度
以我的模型大小(中等,类似于中型变压器)的大小,我无法使用 4 个 GPU 获得大于 ~32 的批量大小,这是我可以获得的最大实例,可悲的是,这些确实是旧的 11GB K80,因为 Azure 似乎认为 Google 甚至不再免费赠送的 GPU 已经足够好......
我有一个需要非常大批量的数据集,因为我必须考虑非常大的不平衡(我也在使用 c 的加权和焦点损失),因此我需要执行 4-8 步的梯度累积来平滑梯度。
我已经阅读了分布式训练循环指南并设法实现它: https ://www.tensorflow.org/beta/tutorials/distribute/training_loops
我还在 TF 2.0 中为自定义训练循环实现了梯度累积,并且tf.keras
:
https ://colab.research.google.com/drive/1yaeRMAwhGkm1voaPp7EtFpSLF33EKhTc