GPU动态增长
import tensorflow as tf
import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
混合精度训练
use_bfloat16 = True
if use_bfloat16:
policy = tf.keras.mixed_precision.experimental.Policy(
'mixed_float16')
tf.keras.mixed_precision.experimental.set_policy(policy)