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分鐘！

成功之後會看到

引用報錯修改