Реализация смешанной точности обучения в TF-Slim

Question

Я пытаюсь использовать смешанную подготовку с tf-slim, чтобы ускорить обучение сетей и использовать тензорные показатели, доступные на моих графических процессорах. Я также хочу использовать несколько сетевых архитектур с предварительно обученными контрольными точками.

Пример обучения смешанной точности и его работы можно найти по адресу https://devblogs.nvidia.com/mixed-precision-resnet-50-tensor-cores/

Основная идея состоит в том, чтобы: 1. Привести входные данные и к fp16 для прямого и обратного прохода. 2. Привести потери обратно к fp32 при корректировке потерь и весов. 3. При использовании потерь для обратного прохода умножить их на шкалу потерь. 4. При обновлении весов делите его на ту же шкалу потерь

Это уменьшает пропускную способность памяти и использует тензорные ядра на графических процессорах Volta и Turing благодаря использованию fp16.

Моя проблема в том, что я не могу понять, куда поместить приведения к fp16 и fp32 с помощью tf-slim.

Чтобы начать обучение, я использую скрипт train_image_classifier.py из models.research.slim

Нужно ли выполнять приведение в файлах определений для сетевых архитектур? Или мне нужно применить изменения в файлах tf.contrib.slim?

Answer 1

Документация NVIDIA по обучению со смешанной точностью дает наглядный пример того, как это сделать с помощью tenorflow.

В Tensorflow реализовано масштабирование потерь в https://www.tensorflow.org/api_docs/python/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)