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

时间:2018-11-03 12:21:11

标签: tensorflow machine-learning conv-neural-network nvidia tf-slim

我正在尝试将混合精度训练与tf-slim结合使用,以加快网络训练速度并利用我GPU上可用的tensorcore。 我还想利用带有预训练检查点的多种网络体系结构。

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

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

这将减少内存带宽,并通过使用fp16来利用Volta和Turing GPU上的Tensor内核。

我的问题是我无法弄清楚使用tf-slim将演员表放到fp16和fp32的位置。

要开始训练,我使用来自models.research.slim的train_image_classifier.py脚本

我是否需要在网络架构的定义文件中进行强制转换? 还是我需要在tf.contrib.slim文件中应用更改?

1 个答案:

答案 0 :(得分: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)