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

Question

在开发我自己的自定义层时tf.keras：我应该如何支持混合精度？混合精度的文档——目前在 Tensorflow 2.2 中被标记为实验性的一个特性——只解释了如何从消费者的角度使用预定义的层（例如tf.keras.layers.Dense一层）。我已经尝试自己猜测并发现了两个 - 可能相关的 - 细节：使用 16 位混合精度时，该dtype属性默认保持不变。float32有一种mixed_precision.get_layer_policy(layer)方法（参见文档）和一种mixed_precision.global_policy()方法（参见文档）可用于检索配置的compute_dtype和variable_dtype。我是否应该使用上述get_layer_policy-method 并将我的变量转换为我层compute_dtype的方法？call(...)（并在创建变量时将variable_dtype我的层build(...)方法传递给？）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(...)自定义层方法的方式吗？

Answer 1

nvidia 的本指南（幻灯片 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)

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