Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to use tf.keras with bfloat16

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')
        )
    )
like image 784
Luke Avatar asked May 13 '19 16:05

Luke


People also ask

Does TensorFlow support bfloat16?

All Cloud TPUs support bfloat16.

Should I use TF keras or keras?

The first important takeaway is that deep learning practitioners using the keras package should start using tf. keras inside TensorFlow 2.0. Not only will you enjoy the added speed and optimization of TensorFlow 2.0, but you'll also receive new feature updates — the latest release of the keras package (v2.

Is TF keras the same as keras?

The difference between tf. keras and keras is the Tensorflow specific enhancement to the framework. keras is an API specification that describes how a Deep Learning framework should implement certain part, related to the model definition and training.

How much faster is mixed precision?

Since the introduction of Tensor Cores in the Volta and Turing architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures.


1 Answers

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()
like image 65
Tensorflow Support Avatar answered Oct 11 '22 12:10

Tensorflow Support