在tf-slim中实现混合精度训练

Question

我正在尝试使用tf-slim的混合精度训练，以加快网络的训练速度，并利用我的GPU上可用的张量核心。我还想使用具有预先训练的检查点的多个网络架构。

有关什么是混合精度训练及其工作原理的示例，请访问https://devblogs.nvidia.com/mixed-precision-resnet-50-tensor-cores/

基本思想是1.将输入转换为正向和反向传递的fp16。2.调整损失和权重时，将值转换回fp32。3.将损失用于反向传递时，将其乘以损失比例4.更新权重时，将其除以相同的损失比例

这减少了内存带宽，并通过使用fp16利用了Volta和图灵GPU上的张量内核。

我的问题是，我不知道如何使用tf-slim将类型转换放到fp16和fp32中。

为了开始培训，我使用了models.research.slim中的train_image_classifier.py脚本

我是否需要在网络体系结构的定义文件中进行转换？或者，我是否需要在tf.contrib.slim文件中应用更改？

Answer 1

NVIDIA的documentation on mixed precision training给出了一个清晰的例子，说明了如何使用tensorflow来实现这一点。

Tensorflow已经在tf.contrib.mixed_precision.LossScaleOptimizer中实现了损失缩放。据我所知，它使用的策略与NVIDIA关于混合精度训练的文档中描述的策略相同。

loss = loss_fn()
opt = tf.AdamOptimizer(learning_rate=...)

# Choose a loss scale manager which decides how to pick the right loss scale 
# throughout the training process.
# Use fixed loss scaling factor
loss_scale_manager = tf.contrib.mixed_precision.FixedLossScaleManager(loss_scale)
# Use dynamic loss scaling factor
loss_scale_manager = tf.contrib.mixed_precision.ExponentialUpdateLossScaleManager(init_loss_scale, 
                                                                                  incr_every_n_steps)

# Wrap the original optimizer in a LossScaleOptimizer.
loss_scale_optimizer = tf.contrib.mixed_precision.LossScaleOptimizer(opt, loss_scale_manager)

# Call minimize() on the loss scale optimizer.
train_op = loss_scale_optimizer.minimize(loss)