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 its 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)

layer = SimpleDense(dtype='mixed_float16')
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
layer = AddRandom(dtype='mixed_float16')
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()