混合精度策略 API

`Policy` 类

tf_keras.mixed_precision.Policy(name)

TF-Keras 层的 dtype 策略。

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

参数

name: 策略名称，决定计算和变量的 dtype。可以是任何 dtype 名称，例如 'float32' 或 'float64'，这将使计算和变量的 dtype 都为该 dtype。也可以是字符串 'mixed_float16' 或 'mixed_bfloat16'，这将使计算 dtype 为 float16 或 bfloat16，而变量 dtype 为 float32。

通常，您只需要在混合精度时与 dtype 策略进行交互，混合精度是指使用 float16 或 bfloat16 进行计算，而使用 float32 作为变量。这就是为什么 API 名称中会出现 mixed_precision 这个术语。可以通过将 '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。这会导致层的计算和输出也以计算 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 是该层的策略的变量 dtype。

如果层的计算 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

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

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

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

您需要额外工作来支持混合精度或 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。您无法将 float32 张量与 float16 张量相加。

[源代码]

`global_policy` 函数

tf_keras.mixed_precision.global_policy()

返回全局 dtype 策略。

全局策略是层的默认 tf.keras.mixed_precision.Policy，如果未将策略传递给层构造函数。如果未通过 keras.mixed_precision.set_global_policy 设置策略，则将返回一个由 tf.keras.backend.floatx()（floatx 默认为 float32）构建的策略。

>>> tf.keras.mixed_precision.global_policy()
<Policy "float32">
>>> tf.keras.layers.Dense(10).dtype_policy  # Defaults to the global policy
<Policy "float32">

如果已通过 tf.compat.v1.disable_v2_behavior() 禁用了 TensorFlow 2 行为，则此函数将返回一个特殊的 "_infer" 策略，该策略会在层首次调用时从第一个输入的 dtype 推断 dtype。此行为与 TensorFlow 1 中的行为匹配。

有关策略的更多信息，请参阅 tf.keras.mixed_precision.Policy。

全局策略。

[源代码]

`set_global_policy` 函数

tf_keras.mixed_precision.set_global_policy(policy)

设置全局 dtype 策略。

全局策略是层的默认 tf.keras.mixed_precision.Policy，如果未将策略传递给层构造函数。

>>> tf.keras.mixed_precision.set_global_policy('mixed_float16')
>>> tf.keras.mixed_precision.global_policy()
<Policy "mixed_float16">
>>> tf.keras.layers.Dense(10).dtype_policy
<Policy "mixed_float16">
>>> # Global policy is not used if a policy
>>> # is directly passed to constructor
>>> tf.keras.layers.Dense(10, dtype='float64').dtype_policy
<Policy "float64">
>>> tf.keras.mixed_precision.set_global_policy('float32')