梯度累加及torch实现

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，因此输出可能会有误差

4. 参考

梯度累加(Gradient Accumulation)