1. 什么是梯度累加
我们在训练神经网络的时候,超参数batch size的大小会对最终的模型效果产生很大的影响。一定条件下,batch size设置的越大,模型就会越稳定。batch size的值通常设置在 8-32 之间,但是当我们做一些计算量需求大的任务(例如语义分割、GAN等)或者输入图片尺寸太大的时候,我们的batch size往往只能设置为2或者4,否则就会出现 “CUDA OUT OF MEMORY” 的不可抗力报错。梯度累加就是在有限的计算资源的条件下变相的扩大batch size
2. 梯度累加的过程
一个神经网络的训练过程通常如下:
- 将前一个batch计算之后的网络梯度清零
- 正向传播,将数据传入网络,得到预测结果
- 根据预测结果与label,计算损失值
- 利用损失进行反向传播,计算参数梯度
- 利用计算的参数梯度更新网络参数
for i, (inputs, labels) in enumerate(trainloader):
optimizer.zero_grad() # 梯度清零
outputs = net(inputs) # 正向传播
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播,计算梯度
optimizer.step() # 更新参数
if (i+1) % evaluation_steps == 0:
evaluate_model()
而梯度累加的过程如下:
- 正向传播,将数据传入网络,得到预测结果
- 根据预测结果与label,计算损失值
- 利用损失进行反向传播,计算参数梯度
- 重复上面步骤,不清空梯度,将梯度累加
- 梯度累加达到固定次数之后,更新参数,然后将梯度清零
for i, (inputs, labels) in enumerate(trainloader):
outputs = net(inputs) # 正向传播
loss = criterion(outputs, labels) # 计算损失函数
loss = loss / accumulation_steps # 损失标准化
loss.backward() # 反向传播,计算梯度
if (i+1) % accumulation_steps == 0:
optimizer.step() # 更新参数
optimizer.zero_grad() # 梯度清零
if (i+1) % evaluation_steps == 0:
evaluate_model()
梯度累加就是每计算一个batch的梯度,不进行清零,而是做梯度的累加,当累加到一定的次数之后,再更新网络参数,然后将梯度清零
3. 实验
def train(trainset, evalset, model, tokenizer, model_dir, lr, epochs, device):
optimizer = AdamW(model.parameters(), lr=lr)
batch_size = 3
gradient_accumulation_steps = 10
total_steps = len(trainset)//gradient_accumulation_steps*epochs
# 每一个epoch中有多少个step可以根据len(DataLoader)计算:total_steps = len(DataLoader) * epoch
# total_steps = (len(trainset)) * epochs
scheduler = get_cosine_schedule_with_warmup(
optimizer, num_warmup_steps=100, num_training_steps=total_steps)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
lr_record = []
for epoch in tqdm(range(epochs), desc="epoch"):
train_loss, steps = 0, 0
for batch in tqdm(trainset, desc="train"):
batch = tuple(input_tensor.to(device) for input_tensor in batch if isinstance(input_tensor, torch.Tensor))
input_ids, label, mc_ids = batch
steps += 1
model.train()
loss, logits = model(input_ids=input_ids, mc_token_ids=mc_ids, labels=label)
# 计算10个batch才算一步,因此loss是10个batch的均值
loss = loss / gradient_accumulation_steps
# loss.backward()
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
train_loss += loss.item()
if steps % gradient_accumulation_steps == 0:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), 5)
# torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
# lr_record.append(scheduler.get_lr()[0])
# if steps % 500 == 0:
# print("step:%d avg_loss:%.3f"%(steps, train_loss/steps))
# plot(lr_record)
eval_res = evaluate(evalset, model, device)
os.makedirs(model_dir, exist_ok=True)
model_path = os.path.join(model_dir, "gpt2clsnews.model%d.ckpt"%epoch)
model.save_pretrained(model_path)
tokenizer.save_pretrained(os.path.join(model_dir,"gpt2clsnews.tokinizer"))
logging.info("checkpoint saved in %s"%model_dir)
结果对比
没有梯度累加 | 带梯度累加 | |
---|---|---|
accuracy | 0.902 | 0.866 |
训练耗时 | 4m3s | 3m25s |
理论上带上梯度累加后效果应该会好,但是此处由于时间原因,只跑了一个epoch,因此输出可能会有误差