如何在自定义 Tensorflow 层中支持混合精度？答案

【问题标题】：How to support mixed precision in custom Tensorflow layers?如何在自定义 Tensorflow 层中支持混合精度？
【发布时间】：2020-10-18 01:00:49
【问题描述】：

在为tf.keras 开发我自己的自定义层时：我应该如何支持混合精度？

documentation of mixed precision - 目前在 Tensorflow 2.2 中被标记为实验性的功能 - 仅说明如何从消费者的角度使用预定义层，例如 tf.keras.layers.Dense。

我已经尝试自己猜测并找到了两个 - 可能相关的 - 详细信息：

使用 16 位混合精度时，dtype 属性默认保持为 float32。
有一个mixed_precision.get_layer_policy(layer) 方法（参见docs）和一个mixed_precision.global_policy() 方法（参见docs）可用于检索配置的compute_dtype 和variable_dtype。

我是否应该使用上述get_layer_policy-方法并将我的变量转换为我层的call(...) 方法中的compute_dtype？（并在创建变量时将我的层build(...) 方法中的variable_dtype 传递给add_weight(...)？）

例如，这里是标准密集神经元层的简单示例实现：

  def call(self, input):
    policy = mixed_precision.get_layer_policy(self)
    bias = tf.cast(self._bias, policy.compute_dtype)
    weights = tf.cast(self._weights, policy.compute_dtype)
    y = tf.nn.bias_add(tf.matmul(input, weights), bias)
    outputs = self._activation(y)
    return outputs

当然，没有人会自己实现这些基本的东西，那只是为了演示。但是，这会是 Tensorflow 团队希望我们实现自定义层的call(...) 方法的方式吗？

【问题讨论】：

标签： python tensorflow tf.keras tensorflow2.x

【解决方案1】：

来自 nvidia 的This 指南（幻灯片 13-14）提到了用于混合精度训练的自定义层。

你必须实现方法cast_input()。在此示例中，当启用混合精度时，图层将转换为 float16：

class CustomBiasLayer(tf.keras.layers.Layer):

 def build(self, _):
   self.v = self.add_weight('v', ())
   self.built = True
  
 def call(self, inputs):
   return inputs + self.v

 def cast_inputs(self, inputs):
   # Casts to float16, the policy's lowest-precision dtype
   return self._mixed_precision_policy.cast_to_lowest(inputs)

我自己没有尝试过，所以如果这对你有用，请告诉我。

【讨论】：