如何使用特斯拉斯带bfloat16

Question

我想弄一个特斯拉斯使用混合精度在TPU上运行的模型。我想知道如何使用bfloat16混合精度构建keras模型。是这样的吗？在

with tf.contrib.tpu.bfloat16_scope():
    inputs = tf.keras.layers.Input(shape=(2,), dtype=tf.bfloat16)
    logits = tf.keras.layers.Dense(2)(inputs)

logits = tf.cast(logits, tf.float32)
model = tf.keras.models.Model(inputs=inputs, outputs=logits)
model.compile(optimizer=tf.keras.optimizers.Adam(.001),
              loss='mean_absolute_error', metrics=[])

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
        model,
        strategy=tf.contrib.tpu.TPUDistributionStrategy(
            tf.contrib.cluster_resolver.TPUClusterResolver(tpu='my_tpu_name')
        )
    )

Answer 1

您可以使用bfloat16Mixed Precision（float16计算和float32变量）来构建Keras模型。在

tf.keras.mixed_precision.experimental.set_policy('infer_float32_vars')

model = tf.keras.Sequential([
    tf.keras.layers.Inputs(input_shape=(2, ), dtype=tf.float16),    
    tf.keras.layers.Lambda(lambda x: tf.cast(x, 'float32')),
    tf.keras.layers.Dense(10)])

model.compile(optimizer=tf.keras.optimizers.Adam(.001),
              loss='mean_absolute_error', metrics=[])

model.fit(.............)

一旦模型被构建和训练，我们可以使用以下步骤保存模型：

^{pr2}$

我们可以使用以下代码加载保存的混合精度Keras模型：

new_model = tf.keras.experimental.load_from_saved_model(path_to_save_model)
new_model.summary()

如果你觉得这个答案有用，请接受这个答案和/或投赞成票。谢谢。在