Python tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite用法及代码示例

通过图形重写启用混合精度。

用法

tf.compat.v1.mixed_precision.enable_mixed_precision_graph_rewrite(
    opt, loss_scale='dynamic'
)

混合精度是在训练模型以提高性能时同时使用 float32 和 float16 数据类型。这是通过图形重写操作和loss-scale 优化器实现的。

在 float16 中执行算术运算利用了专门的处理单元，例如 NVIDIA Tensor Cores，以获得更高的算术吞吐量。但是，由于可表示范围较小，使用 float16 执行整个训练可能会导致梯度下溢，即小的梯度值变为零。相反，当使用兼容的硬件加速器时，仅在 float16 中执行选择算术运算会导致更高的吞吐量和更少的训练时间，同时还可以减少内存使用量，通常不会牺牲模型的准确性。

注意：虽然混合精度重写改变了整个模型中各个层的数据类型，但预计在 float32 中达到的精度相同。如果在动态损失缩放中出现NaN 梯度，则跳过该批次的模型更新。在这种情况下，全局步数不会增加，LossScaleOptimizer 会尝试减小损失缩放值以避免后续迭代中的 NaN 值。这种方法已被证明可以实现与 float32 相同的准确度，并且在大多数情况下，可以实现更好的训练吞吐量。

例子：

model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='softmax'),
])

opt = tf.keras.optimizers.SGD()
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)
model.compile(loss="mse", optimizer=opt)

x_train = np.random.random((1024, 64))
y_train = np.random.random((1024, 64))
model.fit(x_train, y_train)

调用 enable_mixed_precision_graph_rewrite(opt) 在计算梯度之前启用图形重写操作。该函数还返回一个用 LossScaleOptimizer 包的 Optimizer ( opt )。这可以防止在反向传递期间 float16 张量中的下溢。必须将 tf.train.Optimizer 或 tf.keras.optimizers.Optimizer 类型的优化器传递给此函数，然后将其包装以使用损失缩放。

图重写操作将图中某些操作的dtype从float32更改为float16。此重写操作包含或排除了几类操作。以下类别的 Ops 在auto_mixed_precision_lists.h 中类AutoMixedPrecisionLists 下的相应函数中定义：

• ClearList：在数值上没有显著不利影响的操作。例如： ArgMax 和 Floor 。
• AllowList ：被认为在 float16 中执行数字安全的操作，因此总是被转换。例如： Conv2D 。
• DenyList ：在 float16 中执行的操作在数值上不安全，可能会对下游节点产生负面影响。例如： Softmax 。
• GrayList ：被认为在 float16 中执行数字安全的操作，除非来自 DenyList Op 的下游。例如： Add 和 AvgPool 。

使用此函数时，只能通过调用 opt.minimize() 或 opt.compute_gradients() 后跟 opt.apply_gradients() 来计算和应用梯度。不应使用 tf.gradients 或 tf.GradientTape 计算梯度。这是因为返回的优化器将应用损失缩放，而 tf.gradients 或 tf.GradientTape 不会。如果您直接使用 tf.gradients 或 tf.GradientTape ，您的模型可能会由于 float16 下溢问题而无法收敛。

启用 Eager Execution 后，混合精度图形重写仅在 tf.function s 内启用，因为在 tf.function s 之外，没有图形。

对于具有 Tensor 内核的 NVIDIA GPU，作为一般性能指南，如果小于 256，维度(例如批量大小、输入大小、输出大小和通道数)应该是 2 的幂，或者如果大于 256，则可以被 8 整除。信息，请查看 NVIDIA 深度学习性能指南。

目前，混合精度仅在具有计算能力 7.0 及更高版本(Volta、Turing 或更新的架构)的 NVIDIA Tensor Core GPU 上启用。 CPU 和 TPU 上的图形部分不会被图形重写所触及。