AttributeError: module 'tensorflow.keras.mixed_precision' has no attribute 'set_global_policy'

user_8013752 0
Viewed 4230

I need to add mixed precision to my code in order to save some memory. Specifically, I have tried adding mixed precision policy near line 27 in https://github.com/nimRobotics/google-research/blob/master/ravens/train.py, below is the code excerpt

import argparse
import datetime
import os

import numpy as np
from ravens import agents
from ravens import Dataset
import tensorflow as tf

# tf.keras.mixed_precision.set_global_policy('mixed_float16')

# OR

# policy = tf.keras.mixed_precision.Policy('mixed_float16')
# mixed_precision.set_global_policy(policy)

Both the methods result in attribute error as shown below, I'm using Google Colab with TF 2.3.0

Using tf.keras.mixed_precision.set_global_policy('mixed_float16') results in

Traceback (most recent call last):
  File "train.py", line 28, in <module>
    tf.keras.mixed_precision.set_global_policy('mixed_float16')
AttributeError: module 'tensorflow.keras.mixed_precision' has no attribute 'set_global_policy'

Using

policy = tf.keras.mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

results in

Traceback (most recent call last):
  File "train.py", line 29, in <module>
    policy = tf.keras.mixed_precision.Policy('mixed_float16')
AttributeError: module 'tensorflow.keras.mixed_precision' has no attribute 'Policy'

Any help or hints will be highly appreciated!

google-colaboratory keras tensorflow python-3.x python
1 Answers

For tf < 2.4 you should use mixed precision's experimental package.

tf.keras.mixed_precision.experimental.Policy(
    name, loss_scale='auto'
)

For ex, in tf 2.3

policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')
tf.keras.mixed_precision.experimental.set_policy(policy)

and in tf 2.4

tf.keras.mixed_precision.set_global_policy('mixed_float16')

From tf 2.4, this feature is no longer experimental.

user_92157800
Related