深度学习入门（三）一种预测阶段使用任意参数的滑动平均的快捷方式

关于预测阶段从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'))

亲测真的很方便。