tf.train.experimental.enable_mixed_precision_graph_rewrite

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

tf.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([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='softmax'),
])

opt = tf.keras.optimizers.SGD()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
model.compile(loss="mse", optimizer=opt)

x_train = np.random.random((1024, 64))
y_train = np.random.random((1024, 64))
model.fit(x_train, y_train)

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

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

Операция переписывания графа изменяет тип данных некоторых операций в графе с 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, масштабирование потерь не будет применено, что, вероятно, приведёт к тому, что ваша модель не сойдётся из-за проблем утечки float16. Для применения масштабирования потерь с tf.gradients или tf.GradientTape используйте LossScaleOptimizer.get_scaled_loss и LossScaleOptimizer.get_unscaled_gradients. Подробнее см. keras.mixed_precision.experimental.LossScaleOptimizer.

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

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

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

Сравнение со смешанной точностью Keras API

Как эта функция, так и Keras API смешанной точности позволяют использовать смешанную точность в модели. Поэтому можно использовать только один из двух API. Рекомендуется использовать Keras API смешанной точности, так как он более настраиваемый и поддерживает Eager execution. Однако он поддерживает только модели, которые используют слои Keras, в то время как переписывание графа работает с любой моделью, которая использует tf.functionы.

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

Одно из основных преимуществ Keras API заключается в том, что он поддерживает смешанную точность с Eager execution, то есть смешанную точность вне tf.functionов. Переписывание графа будет влиять только на операции внутри tf.functionов, что усложняет отладку в случае возникновения проблем со смешанной точностью. Keras API также более настраиваемый, так как вы можете переопределить любой слой, чтобы он работал в формате float32, передав dtype="float32" в конструктор слоя. Кроме того, вы можете запросить тип данных тензоров в модели, проверив tensor.dtype. При переписывании графа все тензоры, по-видимому, являются float32, так как тип данных изменяется только под капотом.

Основное преимущество переписывания графа (этой функции) заключается в том, что она работает даже если вы не используете слои Keras или какие-либо другие части Keras. Keras API смешанной точности требует моделей, которые используют слои Keras, так как он вставляет преобразования только внутри слоев и моделей Keras. Ещё одним преимуществом является то, что переписывание графа никогда не приводит к TypeError, в то время как Keras API может его вызвать, если вы выполняете определённые операции вне Keras. Например, следующее приведёт к TypeError, если включён Keras API смешанной точности, так как к float16 и float32 тензору будет добавлено значение: tf.keras.layers.Dense(2)(x) + tf.keras.layers.Dense(2, dtype="float32")(x)