Tensorflow在一定迭代次数后停止对某个参数的训练

我的需求是：网络中有很多参数需要训练，其中有一个参数 A 比较特殊，我希望前 100 个 iteration 过程中所有参数同步进行训练，在 100 个 iteration 之后，停止对 A 的更新，即将 A 的值固定，继续训练其他参数。

为了实现这个需求，我用的方法包含以下两部分内容。

1. 网络中设置多个优化器

参考自 StackOverflow 中的一个回答 https://stackoverflow.com/questions/35298326/freeze-some-variables-scopes-in-tensorflow-stop-gradient-vs-passing-variables

该问题与我们的需求还不太一样，他的描述是在 GAN 的训练过程中希望对 G(enerator) 和 D(iscriminator) 进行交替训练，高赞回答提供的思路是设置两个优化器，分别针对不同的参数进行训练。

optimizer = tf.train.AdagradOptimzer(0.01)

first_train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                				"scope/prefix/for/first/vars")
first_train_op = optimizer.minimize(cost, var_list=first_train_vars)

second_train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
								"scope/prefix/for/second/vars")                     
second_train_op = optimizer.minimize(cost, var_list=second_train_vars)

为了从所有参数（很多很多）中单独摘出参数 A，在我的代码中实现的方式是这样的

optim = tf.train.AdamOptimizer(opt.learning_rate, opt.beta1)

first_train_vars = [var for var in tf.trainable_variables() if var.name!='scope/A:0']
first_train_op = optim.minimize(loss, var_list=first_train_vars)

second_train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'scope/A:0')
second_train_op = optim.minimize(loss, var_list=second_train_vars)

别急，这里还没完，因为我们还希望在 100 个 iteration 之后停止对 A 的训练，可以通过设置学习率来实现

2. 设置动态学习率

因为我们希望 100 个 iteration 之后停止对 A 的学习，tensorflow中似乎没有直接实现这种功能的函数（如果您有知道的欢迎告诉我呀）。有一个 tf.stop_gradient() 可以阻止梯度传播，不过这个是在最初建立图的时候就确定了，所以说并不能实现在训练过程中停止对 A 的更新。下面这一部分主要参考了 https://blog.csdn.net/akadiao/article/details/79560731

我们可以通过设置学习率为 0 来实现停止更新，动态修改学习率还是可以实现的。tensorflow提供了很多形式的学习率衰减函数：

• tf.train.piecewise_constant　分段常数衰减
• tf.train.inverse_time_decay　反时限衰减
• tf.train.polynomial_decay　多项式衰减
• tf.train.exponential_decay　指数衰减
• tf.train.natural_exp_decay　自然指数衰减
• tf.train.cosine_decay　余弦衰减
• tf.train.linear_cosine_decay　线性余弦衰减
• tf.train.noisy_linear_cosine_decay　噪声线性余弦衰减

具体的使用可以参考官方文档，这里我就简单采用了分段常数衰减的学习率，设置 100 iteration 之前的学习率为 0.0002，之后的学习率为 0。

boundaries=[100]
learning_rates=[0.0002， 0.]
learning_rate = tf.train.piecewise_constant(
                global_step,
                boundaries=boundaries, 
                values=learning_rates)

second_optim = tf.train.AdamOptimizer(learning_rate, opt.beta1)
second_train_op = second_optim.minimize(loss, var_list=second_train_vars)

因此最终合起来的代码就是

first_train_vars = [var for var in tf.trainable_variables() if var.name!='scope/A:0']
first_optim = tf.train.AdamOptimizer(first_learning_rate, opt.beta1)
first_train_op = first_optim.minimize(loss, var_list=first_train_vars)

boundaries=[100]
learning_rates=[0.0002， 0.]
second_learning_rate = tf.train.piecewise_constant(
                global_step,
                boundaries=boundaries, 
                values=learning_rates)
second_train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'scope/A:0')
second_optim = tf.train.AdamOptimizer(second_learning_rate, opt.beta1)
second_train_op = second_optim.minimize(loss, var_list=second_train_vars)