keras混合精度训练代码

先上配置tf 1.14 keras 2.2.4 GeForce RTX 2080 Ti *1 cuda 10.1
python 3.6
非常非常简单几行代码

import os
os.environ['TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'] = '1'
import tensorflow as tf
opt=tf.compat.v1.train.AdamOptimizer(learning_rate=0.001)
# add a line
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt,loss_scale='dynamic')
my_model1 = Model(inputs=inputs, outputs=y)
my_model1.compile(opt, loss=loss_function1, metrics=metrics1)

可以扩大batch size啦!
速度也肉眼可见的每轮提升10分钟!
成功之后会看到
在这里插入图片描述
引用报错修改

你可能感兴趣的:(keras混合精度训练代码)