|
由于主机显卡只有12g的显存，且只装了一块30系列的卡，因此在跑代码时难免会遇到batch_size不能太大的尴尬，因此可以通过，梯度累计的方式进行优化，来变相扩大batch_size。这样的操作在pytorch中很好实现，但在tf中稍显复杂。
上代码，解释：
def train():
...
...
# 所有可训练参数
trainable_vars = tf.trainable_variables()
# 规定需要训练的参数
vit_trainable_vars = [var for var in trainable_vars if 'VGG' not in var.name] #both generate and vision_transformer #291
print("************vars to train：",len(vit_trainable_vars))
# 在计算图中定义一个操作，为每一个要训练的参数创建一个等shape的全0变量
accum_vars = [tf.Variable(tf.zeros_like(tv.initialized_value()), trainable=False) for tv in vit_trainable_vars]
# 梯度归零操作
zero_ops = [tv.assign(tf.zeros_like(tv)) for tv in accum_vars]
global_step = tf.Variable(1, trainable=False)
# 定义优化操作
with tf.device('/gpu:0'):
with tf.name_scope('train'):
#train_step = tf.train.AdamOptimizer(learning_rate, 0.9, 0.999).minimize(loss, global_step=global_step)
# 优化器
optimizer = tf.train.AdamOptimizer(learning_rate, 0.9, 0.999)
# 计算梯度
grads = optimizer.compute_gradients(loss, vit_trainable_vars)
# 将本次梯度类加
accum_ops = [accum_vars[i].assign_add(gv[0]) for i, gv in enumerate(grads)]
# 优化参数
train_step = optimizer.apply_gradients([(accum_vars[i], gv[1]) for i, gv in enumerate(grads)], global_step=global_step)
...
iter = 0
while True:
...
...
iter += 1
sess.run(accum_ops) #累加两次梯度
if iter % 2 == 0：
...
sess.run(train_step, feed_dict={...}) #优化一次参数
...
sess.run(zero_ops) #将梯度置0
...
这样就完成了计算两次梯度，并累加回传的目的，相当于batch_size扩大了一倍。
值得注意的是，如果我们不规定要保存的参数化，新创建的Variable也会被保存，会导致我们的模型体积变大，因此要规定只保存原来模型的参数。例，我在实际使用中：
var_to_save = [val for val in var if 'Adam' not in val.name and 'Variable_' not in val.name]
saver = tf.train.Saver(var_to_save, max_to_keep=None)
|