How to use tf.keras with bfloat16

Question

I'm trying to get a tf.keras model to run on a TPU using mixed precision. I was wondering how to build the keras model using bfloat16 mixed precision. Is it something like this?

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')
        )
    )

https://cloud.google.com/tpu/docs/bfloat16 can you please this .. — Roshan Bagdiya, May 15 '19 at 17:51
That link doesn't specify how to do it with tf.keras. All the examples are for vanilla tensorflow. — Luke, May 15 '19 at 19:56
You can try that with google colab and see. https://github.com/tensorflow/tensorflow/issues/26759, as of now tf.keras has no bfloat16 support. — ASHu2, May 22 '19 at 02:32
It seemed to say that it has no support for saving a model in hdf5 format. Seems like it might still work to train a model and save in the TF SavedModel format. — Luke, May 22 '19 at 11:04
@TensorflowSupport you're getting that error because I put a fake name in for the TPU. You'll need to put in your own URL there. — Luke, Aug 14 '19 at 17:20

score 4 · Answer 1 · edited Dec 19 '21 at 10:24

You can build the Keras model using bfloat16 Mixed Precision (float16 computations and float32 variables) using the code shown below.

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(.............)

Once the model is Built and Trained, we can Save the model using the below step:

tf.keras.experimental.export_saved_model(model, path_to_save_model)

We can load the Saved Mixed Precision Keras Model using the code below:

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

How to use tf.keras with bfloat16

1 Answers1