随着深度学习模型参数量的增加,现有GPU加载一个深度模型(尤其是预训练模型)后,剩余显存无法容纳很多的训练数据,甚至会仅能容纳一条训练数据。
梯度累积(Gradient Accumulation)是一种不需要额外硬件资源就可以增加批量样本数量(Batch Size)的训练技巧。这是一个通过时间换空间的优化措施,它将多个Batch训练数据的梯度进行累积,在达到指定累积次数后,使用累积梯度统一更新一次模型参数,以达到一个较大Batch Size的模型训练效果。累积梯度等于多个Batch训练数据的梯度的平均值。
TensorFlow 2.0中梯度累积的代码实现:
# 给定累积次数
accum_steps = 5
# 模型训练
for step, x_batch_train in enumerate(x_train):
with tf.GradientTape() as tape:
y_pred = model(x_batch_train,training=True)
loss_value = loss_fn(y_batch_train, y_pred)
# 计算梯度
if step == 0:
accum_grads = [
tf.Variable(tf.zeros_like(tv), trainable=False) for tv in self.trainable_variables
]
accum_grads = [
accum_grads[i].assign_add(grad / accum_batch) for i, grad in enumerate(
self.grad_clipping( # 梯度裁剪
tape.gradient(loss, self.trainable_variables), gard_theta
)
)]
# 参数更新
if (step + 1) % accum_batch == 0:
self.optimizer.apply_gradients(zip(accum_grads, self.trainable_variables))
accum_grads = [tv.assign(tf.zeros_like(tv)) for tv in accum_grads]