tf.keras.mixed_precision.experimental.Policy

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

tf.keras.mixed_precision.experimental.Policy(
    name, loss_scale=USE_DEFAULT
)

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

Для многих моделей политика каждого слоя будет иметь одинаковый тип данных для вычислений и переменных, которые обычно будут float32. В этом случае мы будем ссылаться на единственный тип данных как на тип данных слоя, который можно запросить с помощью свойства tf.keras.layers.Layer.dtype.

Когда используется обучение с смешанной точностью, у большинства слоёв будет тип данных float16 или bfloat16 для вычислений и тип данных float32 для переменных, и поэтому у слоя нет одного единственного типа данных. Когда тип данных переменных не совпадает с типом данных вычислений, переменные будут автоматически приведены к типу данных вычислений, чтобы избежать ошибок типа. В этом случае tf.keras.layers.Layer.dtype ссылается на тип данных переменных, а не на тип данных вычислений. См. руководство по смешанной точности для получения дополнительной информации о том, как использовать смешанную точность.

Некоторые политики также имеют экземпляр tf.mixed_precision.experimental.LossScale, который используется tf.keras.Model для масштабирования потерь. Масштабирование потерь — это техника, используемая с смешанной точностью для предотвращения числового подтекания в градиентах float16. Масштабирование потерь выполняется только моделями в Model.fit, Model.train_on_batch и аналогичных методах. Слои, которые не являются моделями, игнорируют масштаб потерь.

Политики создаются путем передачи строки в конструктор, например tf.keras.mixed_precision.experimental.Policy('float32'). Строка определяет типы данных вычислений и переменных. Она может быть одной из следующих:

• Любое имя типа данных, такое как 'float32' или 'float64'. И тип данных переменных, и тип данных вычислений будут этого типа. По умолчанию масштабирование потерь не выполняется.
• 'mixed_float16' или 'mixed_bfloat16': Тип данных вычислений — float16 или bfloat16, а тип данных переменных — float32. Эти политики используются для обучения с смешанной точностью. С 'mixed_float16' по умолчанию используется динамический масштаб потерь. С 'mixed_bfloat16' масштабирование потерь по умолчанию не выполняется, так как масштабирование потерь не требуется с bfloat16.

Как использовать смешанную точность в модели Keras

Чтобы использовать смешанную точность в модели Keras, можно использовать политику 'mixed_float16' или 'mixed_bfloat16'. tf.keras.mixed_precision.experimental.set_policy может быть использован для установки политики по умолчанию для слоёв, если политика не передана им. Например:

tf.keras.mixed_precision.experimental.set_policy('mixed_float16')
model = tf.keras.models.Sequential([
    tf.keras.layers.Input((100,)),
    # Dense layers use global policy of 'mixed_float16', which does
    # computations in float16 while keeping variables in float32.
    tf.keras.layers.Dense(10),
    tf.keras.layers.Dense(10),
    # Softmax should be done in float32 for numeric stability. We pass
    # dtype='float32' to use float32 instead of the global policy.
    tf.keras.layers.Activation('softmax', dtype='float32')
])

В качестве альтернативы, политика может быть передана в отдельные слои вместо установки глобальной политики с set_policy:

policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
model = tf.keras.models.Sequential([
    tf.keras.layers.Input((100,)),
    tf.keras.layers.Dense(10, dtype=policy),
    tf.keras.layers.Dense(10, dtype=policy),
    # Softmax should be done in float32 for numeric stability.
    tf.keras.layers.Activation('softmax', dtype='float32')
])

Обратите внимание, что политика 'mixed_float16' по умолчанию будет применять масштабирование потерь в Model.fit, Model.train_on_batch и других методах обучения. Если такой метод не используется (например, используется пользовательский цикл обучения) и 'mixed_float16' используется, масштабирование потерь должно быть применено вручную. См. tf.keras.mixed_precision.experimental.LossScaleOptimizer для получения подробностей. Для 'mixed_bfloat16', масштабирование потерь не выполняется, и масштабирование потерь никогда не нужно применять вручную.

См. руководство по смешанной точности для получения дополнительной информации об использовании смешанной точности

Как использовать float64 в модели Keras

Использование float64 аналогично смешанной точности. Глобальная политика может быть установлена на float64 или dtype='float64' может быть передано отдельным слоям. Например, для установки глобальной политики:

tf.keras.mixed_precision.experimental.set_policy('float64')
model = tf.keras.models.Sequential([
    tf.keras.layers.Input((100,)),
    # All layers use global policy of 'float64', which does computations
    # and creates variables in float64.
    tf.keras.layers.Dense(10),
    tf.keras.layers.Dense(10),
    tf.keras.layers.Activation('softmax')
])
# Optionaly set policy back to float32 if any other models use float32
tf.keras.mixed_precision.experimental.set_policy('float32')

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

Слой будет преобразовывать свои входные данные в тип данных вычислений в TensorFlow 2. Например:

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` casts it's inputs to its compute dtype, which is float32, and
# does computations in float32.
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

При написании собственного слоя рекомендуется принимать тензоры только в первом аргументе. Таким образом, все тензоры преобразуются в тип данных вычислений слоя. MyLayer следовательно должен быть написан так:

class MyLayer(tf.keras.layers.Layer):
  # Now, all tensor inputs will be casted.
  def call(self, inputs):
    a, b = inputs
    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.float64

Другие аргументы не преобразуются автоматически по техническим причинам, но это может измениться в будущей незначительной версии.

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

Подкласс слоя может предотвратить автоматическое преобразование своих входных данных, передав autocast=False в конструктор слоя. Например:

class NonAutoCastingLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    kwargs['autocast'] = False
    super(NonAutoCastingLayer, self).__init__(**kwargs)
  def call(self, inp):
    return inp
x = tf.ones((4, 4, 4, 4), dtype='float32')
layer = NonAutoCastingLayer(dtype='float64')
y = layer(x)  # Will not cast inputs to it's compute dtype of float64
y.dtype
tf.float32

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

Тип данных по умолчанию для переменных, созданных с помощью tf.keras.layers.Layer.add_weight, — тип данных переменных политики слоя.

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

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

class MyLayer(tf.keras.layers.Layer):
 def build(self, input_shape):
   self.x = self.add_weight('x')
   self.y = self.add_weight('y', experimental_autocast=False)
policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
layer = MyLayer(dtype=policy)
layer.build((2, 2))
layer.x
<AutoCastVariable 'x:0' shape=() dtype=float32 true_dtype=float32, numpy=...>
layer.y
<tf.Variable 'y:0' shape=() dtype=float32, numpy=...>

Передача experimental_autocast=False полезна для слоёв, которые могут выполнять математические операции внутри переменного типа данных вместо типа данных вычислений. Например, вы можете вычислить статистические данные переменных, такие как среднее значение и дисперсию, в типе данных переменных.

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

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

Например, этот простой плотный слой не требует дополнительных действий для поддержки смешанной точности или float64. Keras автоматически преобразует входные данные и переменные в соответствующий тип данных.

class MyDense(tf.keras.layers.Layer):
  def build(self, input_shape):
    self.kernel = self.add_weight('kernel', (input_shape[-1], 10))
  def call(self, inputs):
    return tf.matmul(inputs, self.kernel)

policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
layer = MyDense(dtype=policy)
x = np.random.rand(10, 10)
y = layer(x)
y.dtype
tf.float16

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

class MyDense(tf.keras.layers.Layer):
  def build(self, input_shape):
    self.kernel = self.add_weight('kernel', (input_shape[-1], 10))
  def call(self, inputs):
    rand = tf.random.normal(shape=inputs.shape, dtype=inputs.dtype)
    return tf.matmul(inputs, self.kernel) + rand

layer = MyDense(dtype=policy)
y = layer(x)
y.dtype
tf.float16

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

Методы

from_config

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

@classmethod
from_config(
    config, custom_objects=None
)

get_config

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

get_config()