关于预测阶段从checkpoint读取模型后如何加载滑动平均的影子变量方法在这篇博客有所介绍:
https://www.cnblogs.com/hrlnw/p/8067214.html
这篇博客介绍方法大致介绍为:
(1)训练阶段保存影子值
(2)预测阶段调用tf.train.ExponentialMovingAverage类的variables_to_restore()方法读取checkpoint中参数。
但是有缺陷的是,variables_to_restore()是返回模型中全部张量列表,并对期望滑动平均的张量取影子,即变量名字后添加‘/ExponentialMovingAverage_1’。但是一般情况下,我们在训练过程中可能只保存部分参数(例如:训练GAN的时候仅保存生成器的参数)
因此我们可以利用tensorflow的集合机制,来保存我们所需要的变量以及其影子。
训练阶段:
(1)保存目标参数到集合 “RAW”
Paras = []
for var in Paras:
tf.add_to_collection('RAW', var)
(2) 获取影子变量并保存到集合"SHADOWS"
EMA = tf.train.ExponentialMovingAverage(decay,steps)
averages_op = EMA.apply(Paras) # apply ema
shadows = [EMA.average(var) for var in Paras]
for var in shadows:
tf.add_to_collection('SHADOWS', var)
(3) 先训练后更新
with tf.control_dependencies([train_opt,averages_op]):
train_opt_ema = tf.no_op(name='train_opt_ema')
(4) 保存
saver = tf.train.Saver(var_list = Paras + shadows)
saver.save(sess, save_path, global_step=steps)
预测阶段:
(1)加载模型并读取集合
meta_graph = tf.train.import_meta_graph(meta_path) # 加载模型
# 读取集合里张量
shadows = tf.get_collection('SHADOWS')
Paras = tf.get_collection('RAW')
(2)影子赋值
assop = [tf.assign(var1, var2) for var1, var2 in zip(Paras,shadows)]
(3)使用滑动平均参数或者原始参数
# 选择EMA参数还是某次迭代参数
if ema:
saver = tf.train.Saver(shadows)
saver.restore(sess, tf.train.latest_checkpoint('./ckpt'))
sess.run(assop)
else:
saver = tf.train.Saver(Paras)
saver.restore(sess, tf.train.latest_checkpoint('./ckpt'))