AttributeError：模块“tensorflow.keras.mixed_ precision”没有属性“set_global_policy”

Question

我需要在代码中添加混合精度以节省一些内存。具体来说，我尝试在https://github.com/nimRobotics/google-research/blob/master/ravens/train.py的第 27 行附近添加混合精度策略，下面是代码摘录

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)

这两种方法都会导致属性错误，如下所示，我使用 Google Colab 和 TF 2.3.0

使用tf.keras.mixed_precision.set_global_policy('mixed_float16')结果于

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'

使用

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

结果是

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'

任何帮助或提示将不胜感激！

Answer 1

因为tf < 2.4您应该使用混合精度的实验包。

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

例如，在tf 2.3

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

并在tf 2.4

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

从 开始tf 2.4，此功能不再是实验性的。