在開發我自己的自定義層時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(...)自定義層方法的方式嗎？
1 回答
小唯快跑啊
TA貢獻1863條經驗 獲得超2個贊
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)
我自己沒有嘗試過，所以如果這對你有用，請告訴我。
添加回答
舉報
0/150
提交
取消