tf.compat.v2.train.experimental.enable_mixed_precision_graph_rewrite

Включить смешанную точность с помощью переписывания графа.

tf.compat.v2.train.experimental.enable_mixed_precision_graph_rewrite(
    opt, loss_scale='dynamic'
)

Смешанная точность — это использование типов данных float32 и float16 при обучении модели для повышения производительности. Это достигается с помощью операции переписывания графа и оптимизатора с масштабированием потерь.

Выполнение арифметических операций в формате float16 использует специализированные вычислительные блоки, такие как NVIDIA Tensor Cores, для значительно большей пропускной способности арифметических операций. Однако из-за меньшего диапазона представимых значений выполнение всего обучения с float16 может привести к потере градиента, то есть малые значения градиента могут стать нулями. Вместо этого выполнение только выбранных арифметических операций в формате float16 приводит к большей пропускной способности и уменьшению времени обучения при использовании совместимых аппаратных ускорителей, а также уменьшению использования памяти, обычно без потери точности модели.

Примечание: Хотя переписывание в смешанной точности изменяет тип данных различных слоёв во всей модели, ожидается достижение той же точности, что и в float32. Если возникает NaN градиент с динамическим масштабированием потерь, обновление модели для этой партии пропускается. В этом случае счётчик глобальных шагов не увеличивается, и LossScaleOptimizer пытается уменьшить значение масштабирования потерь, чтобы избежать NaN значений в последующих итерациях. Этот подход показал достижение такой же точности, что и в float32, и в большинстве случаев лучшей пропускной способности обучения.

Пример:

model = tf.keras.models.Sequential([
  ...
])

opt = tf.keras.optimizers.SGD()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)

model.compile(loss="categorical_crossentropy",
            optimizer=opt,
            metrics=["accuracy"])

model.fit(x_train, y_train,
        batch_size=batch_size,
        epochs=epochs)

Для получения полного примера, демонстрирующего ускорение обучения задачи классификации изображений на CIFAR10, ознакомьтесь с этим блокнотом Colab.

Вызов enable_mixed_precision_graph_rewrite(opt) включает операцию переписывания графа перед вычислением градиентов. Функция дополнительно возвращает Optimizer(opt) заключённую в LossScaleOptimizer. Это предотвращает потерю значимости в тензорах float16 во время обратного прохода.

Оптимизатор типа tf.train.Optimizer или tf.keras.optimizers.Optimizer должен быть передан в эту функцию, который затем будет обернут для использования масштабирования потерь.

http://developer.download.nvidia.com/compute/machine-learning/frameworks/ TF_mixed_precision_training.png" width="500px">

Операция переписывания графа изменяет dtype некоторых операций в графе с float32 на float16. Существует несколько категорий операций, которые либо включаются, либо исключаются с помощью этой операции переписывания. Следующие категории операций определены внутри соответствующих функций в классе AutoMixedPrecisionLists в auto_mixed_precision_lists.h:

• ClearList: Операции, которые не имеют значительных негативных числовых последствий. Например, ArgMax и Floor.
• WhiteList: Операции, которые считаются числово безопасными для выполнения в float16 и поэтому всегда конвертируются. Например, Conv2D.
• BlackList: Операции, которые числово небезопасны для выполнения в float16 и могут негативно повлиять на последующие узлы. Например, Softmax.
• GrayList: Операции, которые считаются числово безопасными для выполнения в float16, за исключением случаев, когда они находятся после операции из списка BlackList. Например, Add и AvgPool.

При использовании этой функции градиенты следует вычислять и применять только с возвращённым оптимизатором, либо вызывая opt.minimize() или opt.compute_gradients() за которым следует opt.apply_gradients(). Градиенты не должны вычисляться с помощью tf.gradients или tf.GradientTape. Это происходит потому, что возвращённый оптимизатор применяет масштабирование потерь, а tf.gradients или tf.GradientTape этого не делают. Если вы непосредственно используете tf.gradients или tf.GradientTape, ваша модель может не сойтись из-за проблем с потерей значимости в float16.

Когда включено немедленное выполнение, переписывание графа в смешанной точности включено только в tf.function, так как вне tf.function нет графа.

Для NVIDIA видеокарт с ядрами Tensor Cores, как общее руководство по производительности, размеры (например, размер пакета, размер входных данных, размер выходных данных и количество каналов) должны быть степенями двойки, если они меньше 256, или в противном случае делиться на 8, если они больше 256. Для получения дополнительной информации ознакомьтесь с руководством по производительности NVIDIA Deep Learning.

В настоящее время смешанная точность включена только на графических процессорах NVIDIA Tensor Core с вычислительной способностью 7.0 и выше (Volta, Turing или более новые архитектуры). Части графа на ЦП и TPU не затрагиваются переписыванием графа. Поддержка TPU появится в скором времени. ЦП не поддерживаются, так как ЦП не выполняют операции float16 быстрее, чем операции float32.