梯度累加及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)