Python tf.keras.mixed_precision.Policy用法及代码示例

Keras 层的 dtype 策略。

用法

tf.keras.mixed_precision.Policy(
    name
)

数据类型策略确定层的计算和可变数据类型。每一层都有一个策略。可以将策略传递给层构造函数的 dtype 参数，或者可以使用 tf.keras.mixed_precision.set_global_policy 设置全局策略。

通常，您只需要在使用混合精度时与 dtype 策略进行交互，即使用 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') 。通常，将 dtype 策略名称传递给层相当于传递相应的策略，因此永远不需要显式构造 Policy 对象。

注意：如果您使用 'mixed_float16' 策略，Model.compile 将使用 tf.keras.mixed_precision.LossScaleOptimizer 自动包装优化器。如果您使用自定义训练循环而不是调用 Model.compile ，则应明确使用 tf.keras.mixed_precision.LossScaleOptimizer 以避免使用 float16 的数字下溢。

层如何使用其策略的计算数据类型

层将其输入转换为其计算 dtype。这会导致层的计算和输出也在计算 dtype 中。例如：

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

如果使用多个输入编写自己的层，您应该在call 中将其他张量显式转换为self.compute_dtype，或者接受第一个参数中的所有张量作为列表。

转换仅在 TensorFlow 2 中发生。如果已调用 tf.compat.v1.disable_v2_behavior()，您可以使用 tf.compat.v1.keras.layers.enable_v2_dtype_behavior() 启用转换行为。

层如何使用其策略的变量 dtype

tf.keras.layers.Layer.add_weight 创建的变量的默认 dtype 是 layer's policy's 变量 dtype。

如果层的计算和变量 dtype 不同，add_weight 将使用称为 AutoCastVariable 的特殊包装器包装浮点变量。 AutoCastVariable 与原始变量相同，只是在 Layer.call 中使用时它会将自身转换为层的计算 dtype。这意味着如果您正在编写一个层，则不必将变量显式转换为层的计算 dtype。例如：

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

图层作者可以通过将 experimental_autocast=False 传递给 add_weight 来防止变量被 AutoCastVariable 包装，这在必须在图层内访问变量的 float32 值时非常有用。

如何编写支持混合精度和 float64 的层。

在大多数情况下，层将自动支持混合精度和 float64，无需任何额外工作，因为基础层会自动转换输入，创建正确类型的变量，并且在混合精度的情况下，使用 AutoCastVariables 包装变量.

您需要额外工作来支持混合精度或 float64 的主要情况是当您创建新张量时，例如使用 tf.ones 或 tf.random.normal ，在这种情况下，您必须创建正确 dtype 的张量。例如，如果您调用 tf.random.normal ，则必须传递计算 dtype，这是输入已转换为的 dtype：

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 的 dtype 默认为 "float32" ，但输入 dtype 是 float16。您不能使用 float16 张量添加 float32 张量。