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

Question

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

可以在https://devblogs.nvidia.com/mixed-precision-resnet-50-tensor-cores/ 找到混合精度训练及其工作原理的示例

基本思想是 1. 将输入和转换为 fp16 用于前向和后向传递 2.调整损失和权重时将值转换回fp32 3.当使用Loss做backward pass时，乘以一个loss scale 4.更新权重时，除以相同的损失尺度

这减少了内存带宽，并通过使用 fp16 来利用 Volta 和 Turing 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)