tf.keras.mixed_precision.Policy

Политика типа данных для слоя Keras.

tf.keras.mixed_precision.Policy(
    name
)

Политика типа данных определяет вычислительный и переменный типы данных слоя. Каждый слой имеет политику. Политики могут быть переданы в аргумент dtype конструкторов слоев, или глобальная политика может быть установлена с помощью tf.keras.mixed_precision.set_global_policy.

Обычно вам нужно взаимодействовать с политиками типа данных только при использовании смешанной точности, которая представляет собой использование float16 или bfloat16 для вычислений и float32 для переменных. Именно поэтому термин mixed_precision появляется в имени API. Смешанная точность может быть включена путем передачи 'mixed_float16' или 'mixed_bfloat16' в tf.keras.mixed_precision.set_global_policy. См. руководство по смешанной точности для получения дополнительной информации о том, как использовать смешанную точность.

tf.keras.mixed_precision.set_global_policy('mixed_float16')
layer1 = tf.keras.layers.Dense(10)
layer1.dtype_policy  # `layer1` will automatically use mixed precision
<Policy "mixed_float16">
# Can optionally override layer to use float32 instead of mixed precision.
layer2 = tf.keras.layers.Dense(10, dtype='float32')
layer2.dtype_policy
<Policy "float32">
# Set policy back to initial float32 for future examples.
tf.keras.mixed_precision.set_global_policy('float32')

В приведенном выше примере передача dtype='float32' слою эквивалентна передаче dtype=tf.keras.mixed_precision.Policy('float32'). В общем случае передача типа данных слою эквивалентна передаче соответствующей политики, поэтому явное создание объекта Policy никогда не требуется.

Примечание: Model.compile автоматически обернёт оптимизатор с tf.keras.mixed_precision.LossScaleOptimizer, если вы используете политику 'mixed_float16'. Если вы используете пользовательскую петлю обучения вместо вызова Model.compile, вы должны явно использовать tf.keras.mixed_precision.LossScaleOptimizer, чтобы избежать численных подпотоков с float16.

Как слой использует вычислительный тип данных своей политики

Слой преобразует свои входные данные в свой вычислительный тип данных. Это приводит к тому, что вычисления и вывод слоя также будут в вычислительном типе данных. Например:

x = tf.ones((4, 4, 4, 4), dtype='float64')
# `layer`'s policy defaults to float32.
layer = tf.keras.layers.Conv2D(filters=4, kernel_size=2)
layer.compute_dtype  # Equivalent to layer.dtype_policy.compute_dtype
'float32'
# `layer` casts it's inputs to its compute dtype and does computations in
# that dtype.
y = layer(x)
y.dtype
tf.float32

Обратите внимание, что базовый класс tf.keras.layers.Layer вставляет преобразования. Если вы создаёте свой собственный слой, вам не нужно вставлять никаких преобразований.

В настоящее время преобразование выполняется только над тензорами в первом аргументе метода call слоя (хотя это, вероятно, будет изменено в будущем). Например:

class MyLayer(tf.keras.layers.Layer):
  # Bug! `b` will not be casted.
  def call(self, a, b):
    return a + 1., b + 1.
a = tf.constant(1., dtype="float32")
b = tf.constant(1., dtype="float32")
layer = MyLayer(dtype="float64")
x, y = layer(a, b)
x.dtype
tf.float64
y.dtype
tf.float32

Если вы создаёте свой собственный слой с несколькими входами, вам следует либо явно преобразовать другие тензоры в self.compute_dtype в call, либо принять все тензоры в первом аргументе в виде списка.

Преобразование выполняется только в TensorFlow 2. Если был вызван tf.compat.v1.disable_v2_behavior(), можно включить поведение преобразования с помощью tf.compat.v1.keras.layers.enable_v2_dtype_behavior().

Как слой использует переменный тип данных своей политики

По умолчанию тип данных переменных, созданных методом tf.keras.layers.Layer.add_weight, равен переменного типу данных политики слоя.

Если вычислительный и переменный типы данных слоя отличаются, add_weight обернёт переменные с плавающей точкой специальной обёрткой, называемой AutoCastVariable. AutoCastVariable идентичен исходной переменной, за исключением того, что он преобразует себя в вычислительный тип данных слоя при использовании в методе Layer.call. Это означает, что если вы пишете слой, вам не нужно явно преобразовывать переменные в вычислительный тип данных слоя. Например:

class SimpleDense(tf.keras.layers.Layer):

  def build(self, input_shape):
    # With mixed precision, self.kernel is a float32 AutoCastVariable
    self.kernel = self.add_weight('kernel', (input_shape[-1], 10))

  def call(self, inputs):
    # With mixed precision, self.kernel will be casted to float16
    return tf.linalg.matmul(inputs, self.kernel)

dtype_policy = tf.keras.mixed_precision.Policy('mixed_float16')
layer = SimpleDense(dtype=dtype_policy)
y = layer(tf.ones((10, 10)))
y.dtype
tf.float16
layer.kernel.dtype
tf.float32

Автор слоя может предотвратить обёртывание переменной с помощью AutoCastVariable, передав experimental_autocast=False в add_weight, что полезно, если значение переменной float32 должно быть обработано внутри слоя.

Как создать слой, поддерживающий смешанную точность и float64

В основном слои будут автоматически поддерживать смешанную точность и float64 без дополнительных действий, поскольку базовый слой автоматически преобразует входные данные, создаёт переменные правильного типа и, в случае смешанной точности, оборачивает переменные в AutoCastVariables.

Основной случай, когда вам нужно добавить дополнительные действия для поддержки смешанной точности или float64, возникает при создании нового тензора, например, с помощью tf.ones или tf.random.normal. В таких случаях вы должны создать тензор правильного типа данных. Например, если вы вызываете tf.random.normal, вы должны передать вычислительный тип данных, который является типом данных, к которому были преобразованы входные данные:

class AddRandom(tf.keras.layers.Layer):

  def call(self, inputs):
    # We must pass `dtype=inputs.dtype`, otherwise a TypeError may
    # occur when adding `inputs` to `rand`.
    rand = tf.random.normal(shape=inputs.shape, dtype=inputs.dtype)
    return inputs + rand

dtype_policy = tf.keras.mixed_precision.Policy('mixed_float16')
layer = AddRandom(dtype=dtype_policy)
y = layer(x)
y.dtype
tf.float16

Если вы не передали dtype=inputs.dtype в tf.random.normal, произошла бы ошибка TypeError. Это связано с тем, что тип данных tf.random.normal по умолчанию равен "float32", но тип данных входных данных — float16. Вы не можете добавить тензор float32 к тензору float16.

Методы

from_config

Просмотреть исходный код

@classmethod
from_config(
    config, custom_objects=None
)

get_config

Просмотреть исходный код

get_config()