如何在自定義 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 -方法並將我的變量轉換為我層的call(...)方法中的compute_dtype ？ （並在創建variable_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(...)方法的方式嗎？

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)

我自己沒有嘗試過，所以如果這對你有用，請告訴我。