iT邦幫忙

2025 iThome 鐵人賽

DAY 10
0
AI & Data

零基礎 AI 入門!從 Wx+b 到熱門模型的完整之路!系列 第 10

【Day 10】用一支「通用訓練器」打天下逐行理解開源Trainer的內容

  • 分享至 

  • xImage
  •  

前言

今天終於進入第 10 天的學習了!雖然我們已經準備好了資料與模型,但卻還不太清楚如何將整個訓練流程封裝成一個乾淨、可重複使用的訓練器。一般來說一個基礎的 Trainer 類別至少需要具備以下功能:完整的訓練/驗證迴圈、最佳模型的保存機制、Early Stopping(提前停止)、學習率排程器,甚至還要能處理 LoRA 的載入與保存。接下來我會分段說明這些設計的考量,以及為什麼必須這樣做。

這支訓練器在解決什麼問題?

簡單來說我們今天的內容就是要把「單個 Epoch 內要完成的工作」與「跨 Epoch 之間需要比較或保存的工作」清楚拆開。在這樣的設計下,訓練器應該提供:

  • 可插拔的資料載入器(train_loadervalid_loader
  • 任意的模型與最佳化器(modeloptimizer
  • 可選擇性啟用的學習率排程器(scheduler
  • Early Stopping 與最佳權重保存(支援 general / LoRA 兩種模式)
  • 訓練與驗證損失的可視化

也就是說這樣的訓練器就是一個樣板化事務的自動化框架,你只需要準備好資料並丟進來,其餘的重複性流程都能交給它處理,大幅減少額外的程式碼負擔。

1. 初始化Trainer

在設計訓練器的時候,第一步就是把所有「可能會變動的東西」都丟進 __init__,這樣之後要改參數或實驗就不用大改程式。

class Trainer:
    def __init__(self, epochs, train_loader, valid_loader, model, optimizer,
                 device=None, scheduler=None, early_stopping=10, save_dir='./checkpoints',
                 load_best_model=False, grad_clip=None, is_lora=False):
        self.epochs = epochs
        self.train_loader = train_loader
        self.valid_loader = valid_loader
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.early_stopping = early_stopping
        self.load_best_model = load_best_model
        self.grad_clip = grad_clip
        self.is_lora = is_lora

        if device is None:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            print('Using device:', self.device)
        else:
            self.device = device

        self.model = model.to(self.device)

        self.save_dir = save_dir
        self.save_name = 'best_model.ckpt'
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)

這裡我特別考慮了三個問題:

  • 裝置選擇:GPU 還是 CPU?
    在深度學習裡,GPU 幾乎是標配,因為它能大幅加速矩陣運算。不過這裡還有一個容易忽略的細節,當 PyTorch 在 GPU 上報錯時訊息往往很模糊,會讓人不知道是哪個層或張量出了問題。

相較之下放在 CPU 上執行,雖然速度慢但錯誤訊息更清晰。因此我在設計 Trainer 的時候,把 device 留作可選參數,如果使用者沒有指定就自動檢查 CUDA 是否可用,能用就跑 GPU,否則退回 CPU。這樣的設計兼顧效能與除錯便利。

LoRA 模型保存與載入的分流

LoRA(Low-Rank Adaptation)是一種高效的微調方法,它的特點是只訓練小部分權重,而不是整個大模型,這也導致它的保存與載入方式和一般模型不太一樣一般模型可以直接 torch.save(model.state_dict()),但 LoRA 模型則需要透過 Hugging Face 的 peft 套件,僅保存 adapter 權重,並在載入時附加到基礎模型上。

為了兼容這兩種情境,我在 Trainer 的初始化裡增加了一個布林參數 is_lora,讓保存與載入流程能夠自動分流,不需要使用者額外操心。

學習率排程器

固定學習率往往不是最佳選擇,若設得過高模型容易震盪甚至無法收斂;設得過低則可能導致訓練進展緩慢,甚至停留在某個次佳解(local minimum)。這正是 學習率排程器(Scheduler)

它存在的理由就是能根據訓練進度動態調整學習率。常見做法包括:StepLR 每隔幾個 epoch 將學習率乘上一個係數、CosineAnnealingLR 讓學習率以餘弦函數方式週期性下降、ReduceLROnPlateau 則在驗證表現停滯時才降低學習率。

這樣的設計能幫助模型在陷入次佳解時,透過調整學習率跳脫卡住的區域,從而提升最終的收斂效果。此外Scheduler 也分為兩種更新頻率:有些在 每個 step 更新,有些則在 每個 epoch 更新。因此在 Trainer 中,我只保留一個 scheduler 參數,讓使用者自由決定要採用哪一種策略。

在訓練器的設計中,我也加入了一些輔助功能來提升實用性。像是 early_stopping,可以設定當模型在驗證集上連續若干個 epoch 沒有進步時就自動停止訓練。接著是 save_dir / save_name,用來統一管理模型檔案的保存位置,避免不同實驗互相覆蓋。load_best_model 則決定在訓練結束後,是否要自動載入表現最佳的權重,省去手動切換的麻煩。最後還有 grad_clip,這是梯度裁剪的設定,用來防止梯度爆炸,特別是在訓練深層模型時非常實用。

單一訓練迴圈的方式

在訓練模型時,使用 self.model.train() 是為了啟用訓練模式,確保像 BatchNormDropout 這些特定模組能正確運作,這兩個模組在訓練和推論時的行為不同:

  • BatchNorm:訓練時根據每個 mini-batch 的數據進行標準化,幫助模型更快收斂;推論時則使用訓練期間累積的統計值,以避免結果不穩定。
  • Dropout:訓練時會隨機關閉部分神經元,降低過擬合風險;推論時則保留所有神經元提供穩定輸出。
    所以我們在訓練與驗證時需要透過 model.train()model.eval() 的切換,讓模型在不同階段用對的方式運作,確保訓練有效、推論穩定。
def train_epoch(self, epoch):
    train_loss = 0
    train_pbar = tqdm(self.train_loader, position=0, leave=True)
    self.model.train()

    for input_datas in train_pbar:
        self.optimizer.zero_grad()
        input_datas = {k: v.to(self.device) for k, v in input_datas.items()}
        outputs = self.model(**input_datas)
        loss = outputs[0]
        loss.backward()

        if self.grad_clip is not None:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)

        self.optimizer.step()

        if self.scheduler is not None:
            self.scheduler.step()

        train_pbar.set_description(f'Train Epoch {epoch}')
        train_pbar.set_postfix({'loss': f'{loss.item():.3f}'})

        train_loss += loss.item()

    return train_loss / len(self.train_loader)

由於我們無法確定模型使用的損失函數或輸入資料格式,因此採用兩個做法來增加彈性:

  1. 模型回傳損失值,方便後續處理;
  2. dict 傳入輸入參數,例如 {k: v.to(self.device) for ...},這樣寫法簡潔,能確保每個張量都正確移到指定設備上,避免遺漏。

此外為了防止梯度爆炸,我們用 clip_grad_norm_ 來限制梯度大小,而學習率則透過排程器動態調整,每個 batch 更新一次,根據設定來決定更新時機(step 或 epoch),讓訓練過程更穩定靈活。

乾淨的驗證流程

在驗證階段,許多訓練時需要的設定其實都不再需要。此時我們只需要模型具備基本的前向傳播功能,因此會透過 self.model.eval() 將模型切換到驗證模式,確保像 Dropout、BatchNorm 這類模組能以推論時的行為運作。

def validate_epoch(self, epoch):
    valid_loss = 0
    valid_pbar = tqdm(self.valid_loader, position=0, leave=True)
    self.model.eval()

    with torch.no_grad():
        for input_datas in valid_pbar:
            input_datas = {k: v.to(self.device) for k, v in input_datas.items()}
            outputs = self.model(**input_datas)
            loss = outputs[0]
            valid_pbar.set_description(f'Valid Epoch {epoch}')
            valid_pbar.set_postfix({'loss': f'{loss.item():.3f}'})
            valid_loss += loss.item()

    return valid_loss / len(self.valid_loader)

而通常為了提升推理效率,還會搭配 torch.no_grad() 停用梯度運算,這麼做可以節省記憶體並加快運算速度,因為驗證過程中並不需要進行反向傳播或更新權重。

主回圈進行Early Stopping 與最佳權重保存

在模型訓練中有幾個實用的技巧能提升效果與效率。首先是 Early Stopping,透過計數器 stop_cnt 追蹤驗證表現是否持續進步。只要連續幾個 epoch 沒有改善,就會提前停止訓練,這不僅能有效避免過擬合,還能節省時間與資源。

接著是昨日提到的最佳權重保存,每當 valid_loss 出現新低時就會立即覆蓋並儲存目前的模型狀態,但要注意在使用 LoRA 時會是要透過 save_pretrained() 儲存,若是一般模型則採用 state_dict() 方式。

def train(self, show_loss=True):
    best_loss = float('inf')
    loss_record = {'train': [], 'valid': []}
    stop_cnt = 0

    for epoch in range(self.epochs):
        train_loss = self.train_epoch(epoch)
        valid_loss = self.validate_epoch(epoch)

        loss_record['train'].append(train_loss)
        loss_record['valid'].append(valid_loss)

        # Save best model
        if valid_loss < best_loss:
            best_loss = valid_loss
            if self.is_lora:
                self.model.save_pretrained(self.save_dir)
            else:
                save_path = os.path.join(self.save_dir, self.save_name)
                torch.save(self.model.state_dict(), save_path)
            print(f'Saving Model With Loss {best_loss:.5f}')
            stop_cnt = 0
        else:
            stop_cnt += 1

        print(f'Train Loss: {train_loss:.5f} | Valid Loss: {valid_loss:.5f} | Best Loss: {best_loss:.5f}\n')

        if stop_cnt == self.early_stopping:
            msg = "Model can't improve, stop training"
            print('-' * (len(msg) + 4))
            print(f'| {msg} |')
            print('-' * (len(msg) + 4))
            break

    if show_loss:
        self.show_training_loss(loss_record)

    if self.load_best_model:
        if self.is_lora:
            from peft import PeftModel
            self.model = PeftModel.from_pretrained(self.model, self.save_dir)
            print(f'Best LoRA model loaded from {self.save_dir}')
        else:
            best_model_path = os.path.join(self.save_dir, self.save_name)
            self.model.load_state_dict(torch.load(best_model_path))
            print(f'Best model loaded from {best_model_path}')

而我們若設定 load_best_model=True,訓練結束後會自動載入表現最好的模型,方便後續的測試操作。

繪製損失曲線看趨勢

這部分和昨天的做法一樣,我們可以透過繪製曲線圖來觀察訓練與驗證的損失變化。如果你是在 Colab 或 Jupyter Notebook 上操作,這樣的圖表呈現已經非常直觀且實用。

def show_training_loss(self, loss_record):
    train_loss, valid_loss = [i for i in loss_record.values()]
    plt.plot(train_loss)
    plt.plot(valid_loss)
    plt.title('Training Loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend(['train', 'valid'], loc='upper left')
    plt.show()

不過注意一點若是在命令列介面(CLI)或遠端伺服器上訓練就建議將圖表儲存成檔案(使用 plt.savefig()),或者輸出到像 TensorBoard 或 Weights & Biases 這類工具中做更進一步的視覺化與紀錄,這樣會較為方便查看時實的動態。
以下是將原本的教學以更完整、條理清晰的文章形式呈現,適合用作部落格、筆記或專案說明文件:


AMP、自動混合精度、梯度累積與多指標支援

在深度學習的訓練過程中,隨著模型規模與資料量的提升,計算效能與記憶體使用變得越來越關鍵。因此我們還可以這麼的優化現有的訓練器。

1. 自動混合精度(AMP: Automatic Mixed Precision)

AMP 是 PyTorch 提供的功能,讓你能在不犧牲模型精度的情況下,自動在 float16 與 float32 之間切換,顯著加速訓練,並節省 GPU 記憶體。我們可以在在 train_epoch 函式中,修改訓練 loop。

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = self.model(**input_datas)
    loss = outputs[0]

scaler.scale(loss).backward()
if self.grad_clip is not None:
    scaler.unscale_(self.optimizer)
    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip)
scaler.step(self.optimizer)
scaler.update()

2. 梯度累積(Gradient Accumulation)

在實際訓練中,如果因 GPU 記憶體限制無法使用較大的 batch size,那麼使用 optimizer.zero_grad() 搭配梯度更新的傳統方式可能會導致訓練不穩定。這時,梯度累積(gradient accumulation)是一種非常實用的技巧。

其核心概念是將多個小 batch 的梯度累加起來,等累積一定步數後再進行一次參數更新。這樣就能模擬大 batch 的效果,同時避免顯存爆炸。

accumulate_steps = 4
for step, input_datas in enumerate(train_pbar, start=1):
    ...
    (loss / accumulate_steps).backward()
    if step % accumulate_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

在這裡,accumulate_steps 設為 4,代表每 4 個 mini-batch 才更新一次權重,而每次的 loss 都會除以 4 以保持梯度的 scale 一致。

3. 多指標回傳與模型保存策略

在驗證階段,我們當然可以根據任務需求靈活地更換評估指標。實際上你可能希望同時計算多個指標(如 Accuracy、F1-score、BLEU 等),以更全面地觀察模型表現。而一個實用的做法是在 validate_epoch() 中建立一個包含多項指標的 metrics 字典。

def validate_epoch(self, ...):
    metric_dict = {
        "accuracy": ...,
        "f1": ...,
        "loss": ...
    }
    return metric_dict

接著在訓練主流程中,你可以指定其中某個主指標作為儲存模型的依據

val_metrics = self.validate_epoch(...)
main_metric = val_metrics["f1"]
if main_metric > self.best_score:
    self.best_score = main_metric
    self.save_model()

雖然 AMP、梯度累積與多指標支援等技巧能有效提升訓練效率,也確實讓模型開發流程更加穩定且具可擴展性,但對我而言,這類設計其實也帶來了更多超參數的調整成本,進一步增加了實作的複雜度。

因此在實務上,我並不常主動採用這些擴充技巧,除非有明確的需求,例如希望縮短訓練時間,或在資源受限下提升表現。這類情境下,再引入這些優化手段,通常能顯著改善整體專案的品質與效率。

下集預告

明天我會帶你實作如何使用 nn.Linear() 建立像 RNN 這樣的神經網路,並在後續的所有章節都會透過我們自定義的 Trainer 啟動訓練流程。這個 Trainer 的設計其實與後續會使用的 Hugging Face 訓練架構是相容的,因此非常值得理解怎麼撰寫。未來你也可能會擴充資料增強、分布式訓練、AMP/Apex、自定義評估指標等等。但只要有了這個 Trainer路就算鋪好了剩下的,就是你在這基礎上,慢慢搭建屬於自己的訓練方式。


上一篇
【Day 9】60% 準確率只是起點用 CNN 與 CIFAR-10 探索深度學習優化之路
系列文
零基礎 AI 入門!從 Wx+b 到熱門模型的完整之路!10
圖片
  熱門推薦
圖片
{{ item.channelVendor }} | {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言