返回文章列表

深度學習模型訓練技巧:學習率與梯度裁剪

本文探討深度學習模型訓練中學習率調整與梯度裁剪技術,包含學習率暖啟動、餘弦衰減與梯度裁剪的原理、實作程式碼範例與說明,並介紹低秩適應(LoRA)技術,以及如何準備資料集、載入預訓練模型、建立 PyTorch 資料集和資料載入器。

機器學習 深度學習

深度學習模型訓練的效率與穩定性,仰賴學習率的調整和梯度裁剪等技術。學習率控制模型引數更新的幅度,梯度裁剪則避免梯度爆炸問題,兩者對於大語言模型(LLM)的訓練尤為關鍵。本文除了探討學習率暖啟動與餘弦衰減的原理及實作,也詳細說明梯度裁剪的計算方式與程式碼範例,更進一步介紹了低秩適應(LoRA)技術,這是一種引數高效的微調方法,可以減少儲存需求並提高可擴充套件性。文章最後也示範瞭如何準備資料集、載入預訓練模型、建立 PyTorch 資料集和資料載入器,為讀者提供實務操作的參考。

強化訓練迴圈的技術:學習率調整與梯度裁剪

在深度學習模型訓練中,特別是在大語言模型(LLM)的訓練過程中,學習率的調整和梯度裁剪是兩個至關重要的技術。本文將探討這兩種技術的原理、實作方法及其在模型訓練中的重要性。

學習率暖啟動與餘弦衰減

學習率暖啟動(Warmup)是一種在訓練初期逐漸增加學習率的技術,目的是避免模型在訓練初期因學習率過高或過低而導致的不穩定。接著,餘弦衰減(Cosine Decay)技術在暖啟動後逐漸降低學習率,以半個餘弦週期的軌跡下降至接近零。

程式碼實作:

import math

min_lr = 0.1 * initial_lr
track_lrs = []
lr_increment = (peak_lr - initial_lr) / warmup_steps
global_step = -1

for epoch in range(n_epochs):
    for input_batch, target_batch in train_loader:
        optimizer.zero_grad()
        global_step += 1
        
        if global_step < warmup_steps:
            lr = initial_lr + global_step * lr_increment
        else:
            progress = ((global_step - warmup_steps) / 
                        (total_training_steps - warmup_steps))
            lr = min_lr + (peak_lr - min_lr) * 0.5 * (
                1 + math.cos(math.pi * progress)
            )
        
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr
        
        track_lrs.append(optimizer.param_groups[0]["lr"])

內容解密:

  1. 學習率初始化:根據 initial_lrpeak_lr 設定初始學習率和峰值學習率。
  2. 暖啟動階段:在前 warmup_steps 步中,學習率線性增加至 peak_lr
  3. 餘弦衰減階段:超出 warmup_steps 後,學習率依照餘弦函式逐漸下降至 min_lr
  4. 學習率更新:每一步更新最佳化器的學習率,並記錄當前學習率。

梯度裁剪

梯度裁剪(Gradient Clipping)是一種透過限制梯度的範數(norm)來穩定模型訓練的技術。當梯度的範數超過設定的閾值時,將梯度按比例縮放至該閾值內。

L2 範數計算範例:

對於一個向量 ( v = [v_1, v_2, …, v_n] ),其 L2 範數定義為: [ ||v||_2 = \sqrt{v_1^2 + v_2^2 + … + v_n^2} ]

若有梯度矩陣: [ G = \begin{bmatrix} 3 & 4 \ 0 & 0 \end{bmatrix} ] 其 L2 範數為: [ ||G||_2 = \sqrt{3^2 + 4^2 + 0^2 + 0^2} = 5 ]

若設定的 max_norm 為 1,則縮放後的梯度矩陣 ( G’ ) 為: [ G’ = G \times \frac{1}{5} = \begin{bmatrix} 0.6 & 0.8 \ 0 & 0 \end{bmatrix} ]

程式碼實作:

from chapter05 import calc_loss_batch

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)

loss = calc_loss_batch(input_batch, target_batch, model, device)
loss.backward()

def find_highest_gradient(model):
    max_grad = None
    for param in model.parameters():
        if param.grad is not None:
            grad_values = param.grad.data.flatten()
            max_grad_param = grad_values.max()
            if max_grad is None or max_grad_param > max_grad:
                max_grad = max_grad_param
    return max_grad

print(find_highest_gradient(model))

內容解密:

  1. 模型初始化:建立模型並將其移至指定裝置。
  2. 損失計算:計算一個批次的損失並進行反向傳播。
  3. find_highest_gradient 函式:遍歷模型引數,找出最大的梯度值。
  4. 梯度檢查:輸出模型中最大的梯度值,用於檢查梯度是否過大。

D.4 修改後的訓練函式

前面的程式碼識別出的最大梯度值為 tensor(0.0411)

現在,讓我們套用梯度裁剪(gradient clipping)並觀察它如何影響最大的梯度值:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
print(find_highest_gradient(model))

套用最大範數(max norm)為1的梯度裁剪後,最大的梯度值明顯變小: tensor(0.0185)

內容解密:

  1. torch.nn.utils.clip_grad_norm_ 函式用於限制梯度的最大範數,防止梯度爆炸。
  2. max_norm=1.0 設定了梯度的最大範數為1.0,超過此值的梯度將被縮放。
  3. 梯度裁剪後的最大梯度值(0.0185)明顯小於之前的最大值(0.0411),顯示梯度裁剪的有效性。

修改後的訓練函式

最後,我們改進了 train_model_simple 訓練函式(參見第5章),加入了三個新概念:線性預熱(linear warmup)、餘弦衰減(cosine decay)和梯度裁剪。這些方法共同幫助穩定大語言模型(LLM)的訓練。

以下是修改後的程式碼,與 train_model_simple 的差異以註解標示:

from chapter05 import evaluate_model, generate_and_print_sample

def train_model(model, train_loader, val_loader, optimizer, device,
                n_epochs, eval_freq, eval_iter, start_context, tokenizer,
                warmup_steps, initial_lr=3e-05, min_lr=1e-6):
    train_losses, val_losses, track_tokens_seen, track_lrs = [], [], [], []
    tokens_seen, global_step = 0, -1
    peak_lr = optimizer.param_groups[0]["lr"]
    total_training_steps = len(train_loader) * n_epochs
    lr_increment = (peak_lr - initial_lr) / warmup_steps
    
    for epoch in range(n_epochs):
        model.train()
        for input_batch, target_batch in train_loader:
            optimizer.zero_grad()
            global_step += 1
            
            if global_step < warmup_steps:
                lr = initial_lr + global_step * lr_increment
            else:
                progress = ((global_step - warmup_steps) /
                            (total_training_steps - warmup_steps))
                lr = min_lr + (peak_lr - min_lr) * 0.5 * (
                    1 + math.cos(math.pi * progress))
            
            # 調整學習率
            for param_group in optimizer.param_groups:
                param_group["lr"] = lr
            track_lrs.append(lr)
            
            loss = calc_loss_batch(input_batch, target_batch, model, device)
            loss.backward()
            
            if global_step > warmup_steps:
                # 套用梯度裁剪
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            optimizer.step()
            tokens_seen += input_batch.numel()
            
            if global_step % eval_freq == 0:
                train_loss, val_loss = evaluate_model(
                    model, train_loader, val_loader,
                    device, eval_iter
                )
                train_losses.append(train_loss)
                val_losses.append(val_loss)
                track_tokens_seen.append(tokens_seen)
                print(f"Ep {epoch+1} (Iter {global_step:06d}): "
                      f"Train loss {train_loss:.3f}, "
                      f"Val loss {val_loss:.3f}")
                generate_and_print_sample(
                    model, tokenizer, device, start_context
                )
    
    return train_losses, val_losses, track_tokens_seen, track_lrs

內容解密:

  1. train_model 函式結合了線性預熱、餘弦衰減和梯度裁剪,以穩定LLM的訓練。
  2. 在預熱階段,學習率線性增加;之後,學習率按照餘弦衰減策略調整。
  3. 梯度裁剪在預熱階段後套用,以防止梯度爆炸。
  4. 訓練過程中,定期評估模型在訓練集和驗證集上的損失,並生成樣本文字。

訓練模型

定義 train_model 函式後,我們可以類別似於使用 train_model_simple 方法來訓練模型:

import tiktoken

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
peak_lr = 5e-4
optimizer = torch.optim.AdamW(model.parameters(), weight_decay=0.1)
tokenizer = tiktoken.get_encoding("gpt2")
n_epochs = 15

train_losses, val_losses, tokens_seen, lrs = train_model(
    model, train_loader, val_loader, optimizer, device, n_epochs=n_epochs,
    eval_freq=5, eval_iter=1, start_context="Every effort moves you",
    tokenizer=tokenizer, warmup_steps=warmup_steps,
    initial_lr=1e-5, min_lr=1e-5
)

內容解密:

  1. 使用 train_model 函式訓練模型,設定了相關的超引數,如 n_epochseval_freqwarmup_steps
  2. 訓練過程中,模型的訓練損失和驗證損失會被記錄和輸出。
  3. 鼓勵讀者使用更大的文字資料集進行訓練,並比較使用 train_modeltrain_model_simple 的結果。

訓練結果

在MacBook Air或類別似筆記型電腦上,訓練大約需要5分鐘,並輸出以下結果:

Ep 1 (Iter 000000): Train loss 10.934, Val loss 10.939
Ep 1 (Iter 000005): Train loss 9.151, Val loss 9.461
Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,

...

Ep 15 (Iter 000130): Train loss 0.041, Val loss 6.915
Every effort moves you?" "Yes--quite insensible to the irony. She wanted him vindicated--and by me!" He laughed again, and threw back his head to look up at the sketch of the donkey. "There were days when I

與預訓練類別似,由於資料集較小且多次迭代,模型在幾個epoch後開始過擬合。儘管如此,我們可以看到函式正在工作,因為它最小化了訓練集損失。

低秩適應(LoRA)技術介紹

低秩適應(LoRA)是一種由 Hu 等人提出的引數高效微調方法(https://arxiv.org/abs/2106.09685)。LoRA 透過學習兩個遠小於原始權重矩陣 $W$ 的矩陣 $A$ 和 $B$ 來近似權重更新 $\Delta W$,其中 $\Delta W \approx AB$,$A$ 和 $B$ 之間進行矩陣乘法運算。

LoRA 的工作原理

使用 LoRA,我們可以將權重更新重新表述為: [ h = Wx + \Delta Wx = Wx + ABx ] 其中 $x$ 是輸入資料,$W$ 是預訓練權重,$\Delta W$ 是權重更新,$A$ 和 $B$ 是 LoRA 矩陣。

圖 E.1 直觀地展示了全微調和 LoRA 的權重更新公式之間的比較。

LoRA 的優勢

LoRA 的優勢在於它能夠在保持預訓練模型權重不變的情況下,透過動態應用 LoRA 矩陣來實作模型定製。這大大減少了儲存需求並提高了可擴充套件性,因為只需要調整和儲存較小的 LoRA 矩陣。

準備資料集

在將 LoRA 應用於垃圾郵件分類別示例之前,我們需要載入資料集和預訓練模型。以下是資料準備的程式碼:

from pathlib import Path
import pandas as pd
from ch06 import (
    download_and_unzip_spam_data,
    create_balanced_dataset,
    random_split
)

url = "https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"
zip_path = "sms_spam_collection.zip"
extracted_path = "sms_spam_collection"
data_file_path = Path(extracted_path) / "SMSSpamCollection.tsv"

download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path)
df = pd.read_csv(data_file_path, sep="\t", header=None, names=["Label", "Text"])
balanced_df = create_balanced_dataset(df)
balanced_df["Label"] = balanced_df["Label"].map({"ham": 0, "spam": 1})
train_df, validation_df, test_df = random_split(balanced_df, 0.7, 0.1)
train_df.to_csv("train.csv", index=None)
validation_df.to_csv("validation.csv", index=None)
test_df.to_csv("test.csv", index=None)

內容解密:

這段程式碼下載並解壓縮垃圾郵件資料集,然後建立平衡的資料集,並將其分為訓練集、驗證集和測試集,最後儲存為 CSV 檔案。

載入預訓練模型

接下來,我們載入預訓練的 GPT 模型:

from gpt_download import download_and_load_gpt2
from chapter04 import GPTModel
from chapter05 import load_weights_into_gpt

CHOOSE_MODEL = "gpt2-small (124M)"
INPUT_PROMPT = "Every effort moves"
BASE_CONFIG = {
    "vocab_size": 50257,
    "context_length": 1024,
    "drop_rate": 0.0,
    "qkv_bias": True
}

內容解密:

這段程式碼指定了要使用的預訓練 GPT 模型的組態,包括詞彙大小、上下文長度、丟棄率和查詢鍵值偏差等引數。

建立 PyTorch 資料集物件

import torch
from torch.utils.data import Dataset
import tiktoken
from chapter06 import SpamDataset

tokenizer = tiktoken.get_encoding("gpt2")
train_dataset = SpamDataset("train.csv", max_length=None, tokenizer=tokenizer)
val_dataset = SpamDataset("validation.csv", max_length=train_dataset.max_length, tokenizer=tokenizer)
test_dataset = SpamDataset("test.csv", max_length=train_dataset.max_length, tokenizer=tokenizer)

內容解密:

這段程式碼建立了 PyTorch 資料集物件,分別用於訓練、驗證和測試。SpamDataset 類別負責載入資料並進行必要的預處理。

建立 PyTorch 資料載入器

from torch.utils.data import DataLoader

num_workers = 0
batch_size = 8
torch.manual_seed(123)

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, num_workers=num_workers, drop_last=False)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, num_workers=num_workers, drop_last=False)

內容解密:

這段程式碼建立了 PyTorch 資料載入器,用於批次載入資料。DataLoader 類別提供了靈活的資料載入方式,可以設定批次大小、工作執行緒數等引數。

驗證資料載入器

print("Train loader:")
for input_batch, target_batch in train_loader:
    pass
print("Input batch dimensions:", input_batch.shape)
print("Label batch dimensions", target_batch.shape)

內容解密:

這段程式碼驗證了資料載入器的正確性,列印出了輸入批次和標籤批次的維度。

列印批次數量

print(f"{len(train_loader)} training batches")
print(f"{len(val_loader)} validation batches")
print(f"{len(test_loader)} test batches")

內容解密:

這段程式碼列印出了訓練、驗證和測試資料載入器中的批次數量。