如何在自定义 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)我自己没有尝试过，所以如果这对你有用，请告诉我。