tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite

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

tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite(
    opt, loss_scale='dynamic'
)

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

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

Примечание: Хотя переработка смешанной точности изменяет тип данных различных слоев модели, ожидается достижение той же точности, что и в float32. Если произойдет подтекание градиента при динамическом масштабировании потерь, обновление модели для этой группы пропускается. В этом случае счетчик глобальных шагов не увеличивается, и происходит попытка уменьшить значение масштабирования потерь, чтобы избежать 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.train.Optimizer или tf.keras.optimizers.Optimizer должен быть передан в эту функцию, которая затем будет обернута для использования масштабирования потерь.

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

• ClearList: Операции, которые не оказывают существенного негативного влияния на точность. Например, ArgMax и Floor.
• AllowList: Операции, которые считаются безопасными для выполнения в формате float16, и поэтому всегда преобразуются. Например, Conv2D.
• DenyList: Операции, которые считаются небезопасными для выполнения в формате float16 и могут отрицательно сказаться на последующих узлах. Например, Softmax.
• GrayList: Операции, которые считаются безопасными для выполнения в формате float16, за исключением случаев, когда они находятся ниже по потоку от запрещенной операции. Например, 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 не затрагиваются переработкой графа.