梯度累計講解-支援更大的batch

海_纳百川發表於2024-08-07
梯度BAT

在對比學習（Contrastive Learning）中，梯度累計（Gradient Accumulation）是一種技術，用於在記憶體有限的情況下實現大批次（Large Batch）訓練。這個操作透過將多個小批次的梯度累加起來，再進行一次權重更新，從而模擬大批次訓練的效果。

以下是梯度累計的基本操作步驟：

  1. 初始化：在訓練開始時，初始化模型引數和最佳化器。

  2. 設定累計步數：定義一個引數 accumulation_steps，表示要累積多少個小批次的梯度後進行一次權重更新。

  3. 累積梯度：

    • 對每一個小批次（Mini-batch）資料，進行前向傳播，計算損失函式。
    • 進行反向傳播，計算梯度，但不進行權重更新。此時，梯度會累加到當前的梯度快取中。
    • 每處理一個小批次的資料，增加一個步數計數器 step。

  4. 權重更新：

    • 當 step 達到 accumulation_steps，進行一次權重更新。此時，最佳化器會使用累積的梯度進行引數更新。
    • 重置步數計數器 step 和梯度快取。

  5. 重複：重複上述步驟，直到完成所有的訓練資料。

以下是一個使用 PyTorch 的虛擬碼示例，展示瞭如何實現梯度累計：

import torch
import torch.nn as nn
import torch.optim as optim

# 模型和最佳化器
model = YourModel()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss()

# 設定累計步數
accumulation_steps = 4

# 訓練迴圈
for epoch in range(num_epochs):
    optimizer.zero_grad()  # 初始化梯度
    for i, (inputs, labels) in enumerate(dataloader):
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()  # 反向傳播，計算梯度
        
        # 累積梯度，每 accumulation_steps 更新一次權重
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()  # 更新權重
            optimizer.zero_grad()  # 清空累計的梯度

# 最後的更新，確保剩餘的梯度也被更新
if (i + 1) % accumulation_steps != 0:
    optimizer.step()
    optimizer.zero_grad()

在對比學習的具體應用中，例如 SimCLR 或 MoCo，梯度累計同樣適用，特別是在記憶體受限的環境中。透過梯度累計，可以有效地提高模型訓練的穩定性和收斂速度，同時模擬大批次訓練的效果。

使用大批次訓練通常可以提高訓練的穩定性和效率，但如果批次大小過大，可能會超過顯示卡的視訊記憶體限制，導致記憶體溢位（out-of-memory, OOM）錯誤。為了在有限的視訊記憶體中實現大批次訓練，梯度累計技術應運而生。

透過梯度累計，可以將多個小批次的資料逐個處理，每次計算的梯度累加起來，最後在累積了一定次數後進行一次引數更新。這種方法相當於模擬大批次訓練，同時避免視訊記憶體不足的問題。

例如，如果你想要一個等效的批次大小為 256 的訓練，但由於視訊記憶體限制只能使用批次大小為 64 的小批次，可以設定梯度累計步數為 4（256 / 64 = 4）。這樣每處理 4 個小批次的資料才進行一次權重更新，效果上等同於使用批次大小為 256 的訓練。

相關文章