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, для значительного увеличения пропускной способности арифметических вычислений. Однако из-за меньшего диапазона представимых значений выполнение всего процесса обучения с 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.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, если только они не следуют за операцией из DenyList. Например, Add и AvgPool.

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

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

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

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