Day 11: Training Loop, Optimizer, Scheduler

Mục tiêu của ngày học

Sau bài này, bạn cần làm được 8 việc:

1. Viết được training loop PyTorch đúng thứ tự: load batch, forward, loss, zero_grad, backward, gradient clipping, optimizer.step, scheduler.
2. Tách rõ train loop, validation loop và test loop bằng model.train(), model.eval() và torch.no_grad() hoặc torch.inference_mode().
3. Giải thích được vai trò của optimizer, learning rate, weight decay và scheduler.
4. Biết khi nào chọn SGD, Adam, AdamW, ReduceLROnPlateau, OneCycleLR, warmup hoặc cosine decay.
5. Thêm được early stopping, checkpoint best model và resume metadata.
6. Log được loss, metric, learning rate, gradient norm, epoch time và config.
7. Hiểu performance trade-off của batch size, num_workers, pin_memory, mixed precision và checkpoint IO.
8. Trả lời được: "Dùng được trong production không? Nếu có thì cần điều kiện gì?"

Bài này nối logic từ Day 9 và Day 10 như thế nào?

Day 9 giúp bạn thấy neural network từ bên trong: forward pass, loss, gradient và manual update bằng NumPy. Day 10 thay phần manual bằng PyTorch primitive: Tensor, autograd, nn.Module, Dataset, DataLoader và device management.

Day 11 là bước tiếp theo: biến các primitive đó thành một training job có thể kiểm soát được. Đây là phần khác biệt giữa "chạy được một notebook" và "có một pipeline training có thể debug, repeat, rollback và tối ưu chi phí".

Day 9  -> hiểu toán và gradient
Day 10 -> dùng PyTorch primitive đúng cách
Day 11 -> tổ chức training job gần production
Day 16 -> fine-tune PhoBERT/BERT classifier

Cách học đề xuất trong 2.5-3 giờ

TL;DR

Training loop là runtime engine của deep learning. Một batch đi qua model, model tạo prediction, loss đo sai số, autograd tính gradient, optimizer cập nhật weights. Scheduler điều chỉnh learning rate theo thời gian hoặc theo validation metric.

Loop tối thiểu:

model.train()
for xb, yb in train_loader:
    xb = xb.to(device)
    yb = yb.to(device)

    optimizer.zero_grad(set_to_none=True)
    logits = model(xb)
    loss = loss_fn(logits, yb)
    loss.backward()
    optimizer.step()

Loop gần production cần thêm:

• Seed, config, device và dataset split có thể tái lập.
• Dataset/DataLoader tách khỏi model.
• Train/eval mode đúng.
• Validation metric để chọn checkpoint tốt nhất.
• Gradient clipping để giảm rủi ro exploding gradient.
• Scheduler để kiểm soát learning rate.
• Early stopping để tiết kiệm compute và giảm overfitting.
• Logging đủ để debug loss spike, NaN, learning rate và thời gian chạy.
• Checkpoint có model state, optimizer state, scheduler state, scaler state nếu dùng mixed precision, config và metric.

Mental model cho Senior SE

Điểm quan trọng: training loop không chỉ là code tính toán. Nó là một workflow có dữ liệu, artifact, config, metric, chi phí, khả năng reproduce và khả năng rollback.

Anatomy chuẩn của một epoch

for epoch:
  model.train()
  for train batch:
    move batch to device
    optimizer.zero_grad()
    forward
    loss
    backward
    clip gradient if needed
    optimizer.step()

  model.eval()
  with no_grad or inference_mode:
    for validation batch:
      forward
      loss
      metric

  scheduler.step(...)
  save best checkpoint if validation improves
  early stop if validation stops improving

Các điểm dễ sai:

• Quên optimizer.zero_grad() làm gradient cộng dồn qua batch ngoài ý muốn.
• Quên model.eval() làm Dropout/BatchNorm hoạt động sai khi validation.
• Quên torch.no_grad() hoặc torch.inference_mode() làm tốn memory trong validation.
• Gọi scheduler sai thời điểm, đặc biệt ReduceLROnPlateau cần validation loss.
• Save checkpoint cuối epoch thay vì checkpoint tốt nhất theo validation metric.
• Chỉ log train loss, không log validation metric nên không biết overfitting.

Best solution theo context

Không có optimizer hoặc scheduler tốt nhất cho mọi bài toán. Best solution phụ thuộc dữ liệu, model, loss surface, hardware, latency/cost budget và mục tiêu metric.

Dùng được trong production không? Nếu có thì cần điều kiện gì?

Có. Training loop PyTorch kiểu trong bài có thể là nền cho production training nếu được đóng gói thành script/job có guardrail đầy đủ. Code trong exercise vẫn là bản học tập nhưng đã chứa các building block quan trọng.

Điều kiện tối thiểu để dùng production:

• Data split cố định, không leakage, có dataset snapshot hoặc data version.
• Config được lưu cùng artifact: model architecture, feature schema, optimizer, scheduler, seed, threshold, metric và package version.
• Training job có log structured: loss, metric, learning rate, gradient norm, epoch time, data size, device, checkpoint path.
• Checkpoint lưu best model theo validation metric, không chỉ epoch cuối.
• Có cơ chế resume hoặc ít nhất rollback từ artifact đã lưu.
• Evaluation trên test set chỉ chạy sau khi chọn model bằng validation set.
• Có monitoring cho NaN, loss spike, class imbalance, data quality và drift sau deploy.
• Checkpoint chỉ load từ nguồn tin cậy; không log raw PII hoặc sample nhạy cảm.
• Có benchmark throughput, GPU/CPU memory, checkpoint IO và training cost.
• Có owner chịu trách nhiệm khi metric production lệch khỏi offline metric.

Không nên coi code training loop trong notebook là production nếu thiếu reproducibility, checkpoint, validation, observability và artifact management.

Deliverable cuối ngày

Bạn nên có 4 output:

1. Một training run hoàn chỉnh chạy được trên CPU hoặc GPU.
2. Một checkpoint best model được lưu theo validation loss.
3. Một bảng log theo epoch có train loss, validation loss, validation F1, learning rate và thời gian chạy.
4. Một đoạn production review trả lời vì sao training job này có thể hoặc chưa thể đưa vào production.

Checklist hoàn thành

• Tôi giải thích được thứ tự zero_grad -> forward -> loss -> backward -> clip -> optimizer.step.
• Tôi biết vì sao train dùng model.train() còn validation/test dùng model.eval().
• Tôi biết khi nào dùng torch.no_grad() hoặc torch.inference_mode().
• Tôi chọn được AdamW hoặc SGD theo context và giải thích trade-off.
• Tôi biết scheduler nào gọi theo epoch, scheduler nào gọi theo validation metric, scheduler nào gọi theo step.
• Tôi thêm được gradient clipping và biết phải unscale gradient trước khi clipping nếu dùng AMP.
• Tôi save được checkpoint gồm model, optimizer, scheduler, scaler, config, epoch và best metric.
• Tôi có early stopping dựa trên validation metric.
• Tôi log được metric đủ để debug một training run.
• Tôi nêu được điều kiện để training loop PyTorch dùng được trong production.

Tài liệu

1. Vì sao Day 11 quan trọng?

Một model deep learning không chỉ cần architecture đúng. Model còn cần training process đúng. Cùng một nn.Module, nếu training loop sai, kết quả có thể rất tệ dù code không crash:

• Loss không giảm vì learning rate quá cao hoặc quá thấp.
• Validation metric tốt giả tạo vì data leakage.
• Model overfit vì không có validation/early stopping.
• Metric không reproduce được vì seed và split thay đổi.
• Checkpoint không resume được vì chỉ lưu model weights, không lưu optimizer/scheduler state.
• GPU dùng kém vì data loading bottleneck hoặc batch size quá nhỏ.
• Training silently sai vì quên model.eval() khi validation.

Với Senior SE, hãy nhìn training như một stateful batch job có artifact. Nó cần config, logging, monitoring, retry, rollback và cost control giống các hệ thống production khác.

2. Training loop anatomy

Một training step chuẩn gồm các bước:

1. Lấy batch từ DataLoader
2. Move input và label sang device
3. Reset gradient cũ
4. Forward pass
5. Tính loss
6. Backward pass
7. Optional: gradient clipping
8. Optimizer update
9. Optional: log metric theo batch hoặc epoch

Code tối thiểu:

model.train()

for xb, yb in train_loader:
    xb = xb.to(device)
    yb = yb.to(device)

    optimizer.zero_grad(set_to_none=True)
    logits = model(xb)
    loss = loss_fn(logits, yb)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

Vì sao zero_grad đứng trước forward?

PyTorch accumulate gradient vào .grad của parameter. Đây là thiết kế có chủ ý vì có trường hợp cần gradient accumulation qua nhiều micro-batch. Nhưng trong training thông thường, mỗi optimizer update chỉ nên dùng gradient của batch hiện tại. Vì vậy cần reset gradient trước khi tính batch mới.

optimizer.zero_grad(set_to_none=True) thường tốt hơn set gradient về zero vì:

• Giảm thao tác ghi memory.
• Có thể tiết kiệm memory trong một số workload.
• Dễ phát hiện parameter không nhận gradient vì .grad vẫn là None.

Trade-off: code cũ hoặc custom logic nào giả định .grad luôn là tensor zero có thể cần sửa.

Vì sao validation phải tách khỏi training?

Training và validation có mục đích khác nhau:

model.train() bật behavior training của layer như Dropout và BatchNorm. model.eval() chuyển chúng sang evaluation behavior. torch.no_grad() hoặc torch.inference_mode() tắt autograd để giảm memory và compute overhead.

Lỗi phổ biến: chỉ dùng torch.no_grad() nhưng quên model.eval(). Khi đó Dropout vẫn random và BatchNorm vẫn dùng training behavior, làm metric validation không ổn định.

3. Loss, metric và threshold

Loss là objective optimizer minimize. Metric là thứ bạn dùng để ra quyết định theo business context. Chúng không nhất thiết giống nhau.

Ví dụ binary classification:

• Loss: BCEWithLogitsLoss.
• Metric: F1, recall, precision, PR-AUC, ROC-AUC hoặc accuracy.
• Business policy: threshold chọn theo cost false positive/false negative.

Nên dùng BCEWithLogitsLoss thay vì Sigmoid + BCELoss vì implementation này ổn định số học hơn: model output raw logits, loss tự xử lý phần sigmoid bên trong.

loss_fn = torch.nn.BCEWithLogitsLoss()
logits = model(xb)
loss = loss_fn(logits, yb)
probs = torch.sigmoid(logits)

Production note:

• Không chọn threshold chỉ vì default 0.5. Hãy chọn theo validation set và business cost.
• Với class imbalance, accuracy có thể gây hiểu nhầm. Ưu tiên precision/recall/F1/PR-AUC.
• Log cả loss và metric. Loss giảm không đảm bảo metric business tăng.

4. Optimizer là gì?

Optimizer nhận gradient và quyết định cách update parameter.

parameter_new = parameter_old + update_rule(gradient, learning_rate, optimizer_state)

Optimizer có state nội bộ. Ví dụ AdamW lưu moving average của gradient và squared gradient. Vì vậy checkpoint production nên lưu cả optimizer.state_dict(), không chỉ model.state_dict().

SGD

SGD update trực tiếp theo gradient:

w = w - lr * gradient

Điểm mạnh:

• Đơn giản, dễ hiểu, ít state.
• Có thể generalize tốt trong một số bài vision/classic deep learning.
• Dễ debug khi học optimization.

Điểm yếu:

• Nhạy với learning rate.
• Converge chậm nếu loss surface phức tạp.
• Thường cần momentum và schedule tốt.

Nên dùng khi:

• Cần baseline tối giản.
• Model/dataset đủ quen thuộc.
• Bạn có thời gian tuning learning rate và momentum.

SGD + momentum

Momentum thêm "quán tính" để giảm zig-zag:

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    weight_decay=1e-4,
)

Trade-off: converge mượt hơn SGD thường, nhưng vẫn cần tuning. Với Transformer fine-tuning, đây thường không phải lựa chọn đầu tiên.

Adam

Adam dùng adaptive learning rate theo từng parameter dựa trên moving average của gradient.

Điểm mạnh:

• Prototype nhanh.
• Ít nhạy hơn SGD với learning rate ban đầu.
• Tốt cho nhiều bài NLP, sparse feature hoặc model khó optimize.

Điểm yếu:

• Có nhiều state hơn, tốn memory hơn.
• Weight decay trong Adam truyền thống không tách biệt tốt như AdamW.
• Không luôn generalize tốt hơn SGD.

AdamW

AdamW là lựa chọn mặc định tốt cho nhiều bài hiện đại, đặc biệt Transformer và fine-tuning. Điểm khác quan trọng là weight decay được decouple khỏi gradient update.

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    weight_decay=0.01,
)

Nên bắt đầu với AdamW khi:

• Fine-tune BERT/PhoBERT hoặc model pretrained.
• Cần baseline mạnh nhanh.
• Không muốn mất nhiều thời gian tuning optimizer ban đầu.

Trade-off:

• Tốn memory hơn SGD vì optimizer state.
• Learning rate vẫn cần chọn cẩn thận.
• Weight decay không nên áp dụng bừa bãi cho mọi parameter trong model lớn; với Transformer production thường tách group để không decay bias và LayerNorm weight.

5. Learning rate là hyperparameter nhạy nhất

Learning rate quyết định độ lớn của mỗi update.

LR quá cao  -> loss dao động, NaN, model không converge
LR quá thấp -> training rất chậm, dễ underfit trong budget cố định
LR hợp lý   -> loss giảm tương đối mượt, validation metric cải thiện

Triệu chứng thường gặp:

Best practice: khi debug, hãy thử overfit một batch nhỏ. Nếu model không thể overfit 32-128 samples, có thể có bug về loss, label, shape, optimizer hoặc data.

6. Scheduler là gì?

Scheduler thay đổi learning rate trong quá trình training. Lý do: LR tốt ở đầu training có thể quá cao ở cuối training.

StepLR

Giảm LR theo chu kỳ cố định.

scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=10,
    gamma=0.5,
)

for epoch in range(num_epochs):
    train_one_epoch(...)
    validate(...)
    scheduler.step()

Nên dùng khi cần baseline dễ explain. Nhược điểm là lịch giảm cứng, không phản ứng với validation metric.

ReduceLROnPlateau

Giảm LR khi metric không cải thiện.

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.5,
    patience=2,
)

for epoch in range(num_epochs):
    train_one_epoch(...)
    val_loss = evaluate(...)
    scheduler.step(val_loss)

Nên dùng cho project vừa và nhỏ khi validation loss là tín hiệu đáng tin. Lưu ý scheduler này nhận metric, không gọi trống như StepLR.

OneCycleLR

Thay đổi LR theo từng step, thường tăng rồi giảm trong một cycle. Cần biết tổng số step.

scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-3,
    epochs=num_epochs,
    steps_per_epoch=len(train_loader),
)

for epoch in range(num_epochs):
    for batch in train_loader:
        train_step(...)
        scheduler.step()

Nên dùng khi muốn train nhanh và đã kiểm soát tốt số step. Gọi sai theo epoch thay vì theo batch sẽ làm schedule sai.

Cosine decay và warmup

Cosine decay giảm LR mượt từ cao xuống thấp. Warmup tăng LR dần ở đầu training. Warmup rất phổ biến khi fine-tune Transformer vì pretrained weights nhạy với update quá lớn ngay từ những step đầu.

Mental model:

warmup:     bảo vệ model khỏi update sốc ở đầu training
decay:      giảm update khi model gần vùng tốt
plateau:    phản ứng khi validation không cải thiện
one-cycle:  tăng tốc training khi đã biết total steps

7. Dropout, weight decay, gradient clipping, early stopping

Dropout

Dropout random tắt một phần activation trong training để model không phụ thuộc quá mức vào vài neuron.

model = torch.nn.Sequential(
    torch.nn.Linear(input_dim, 128),
    torch.nn.ReLU(),
    torch.nn.Dropout(p=0.2),
    torch.nn.Linear(128, 1),
)

Trade-off:

• Giảm overfitting khi model lớn hoặc data ít.
• Có thể làm underfit nếu p quá cao.
• Chỉ active trong model.train(), bị tắt trong model.eval().

Weight decay

Weight decay phạt weight quá lớn, giúp regularization.

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    weight_decay=0.01,
)

Trade-off:

• Giúp giảm overfitting.
• Quá cao làm model khó fit.
• Với model lớn, nên tách parameter group để tránh decay bias/LayerNorm khi cần.

Gradient clipping

Gradient clipping giới hạn norm của gradient để giảm exploding gradient.

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

Khi dùng AMP, gradient đang được scale. Cần unscale trước khi clipping:

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

Trade-off:

• Tăng stability, đặc biệt với RNN, Transformer, loss spike hoặc data noisy.
• Nếu max norm quá thấp, update bị bóp quá mạnh và training chậm.
• Không thay thế việc chọn learning rate đúng.

Early stopping

Early stopping dừng training khi validation metric không cải thiện sau một số epoch.

if val_loss improved:
    save best checkpoint
    stale_epochs = 0
else:
    stale_epochs += 1

if stale_epochs >= patience:
    stop training

Trade-off:

• Tiết kiệm compute và giảm overfitting.
• Có thể dừng quá sớm nếu validation metric noisy.
• Cần min_delta và patience phù hợp.
• Luôn save best checkpoint, không dùng weights của epoch cuối nếu epoch cuối tệ hơn.

8. Mixed precision overview

Mixed precision dùng dtype thấp hơn như FP16 hoặc BF16 cho một số phép toán để giảm memory và tăng throughput trên GPU phù hợp. PyTorch cung cấp torch.amp.autocast và torch.amp.GradScaler.

Pattern CUDA AMP hiện đại:

scaler = torch.amp.GradScaler(device="cuda", enabled=use_amp)

for xb, yb in train_loader:
    optimizer.zero_grad(set_to_none=True)

    with torch.amp.autocast(
        device_type="cuda",
        dtype=torch.float16,
        enabled=use_amp,
    ):
        logits = model(xb)
        loss = loss_fn(logits, yb)

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

Khi nên dùng:

• Training trên NVIDIA GPU có Tensor Cores.
• Model đủ lớn để memory/throughput là bottleneck.
• Bạn có metric validation để xác nhận không mất ổn định.

Khi nên tắt để debug:

• Loss thành NaN hoặc metric bất thường.
• Custom operation chưa ổn định với low precision.
• Chạy CPU hoặc workload quá nhỏ, overhead không đáng.

Production note: AMP là optimization, không phải requirement. Bắt đầu bằng FP32 để xác minh correctness, sau đó benchmark AMP.

9. DataLoader và performance

DataLoader không chỉ để chia batch. Nó ảnh hưởng trực tiếp đến throughput.

Best practice:

• Dataset không nên tự move tensor lên GPU. Training loop quyết định device.
• Train loader nên shuffle. Validation/test loader không cần shuffle.
• Nếu GPU utilization thấp, kiểm tra data loading trước khi tăng model size.
• Với NLP, bottleneck có thể nằm ở tokenization. Cân nhắc cache tokenized dataset cho training lặp lại.

10. Checkpoint đúng cần lưu gì?

Có 2 loại artifact:

Checkpoint training nên có:

payload = {
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "scheduler_state_dict": scheduler.state_dict(),
    "scaler_state_dict": scaler.state_dict(),
    "best_val_loss": best_val_loss,
    "config": config_dict,
    "metrics": validation_metrics,
}
torch.save(payload, checkpoint_path)

Load checkpoint:

checkpoint = torch.load(
    checkpoint_path,
    map_location=device,
    weights_only=False,
)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])

Security note: checkpoint PyTorch có thể dùng pickle bên dưới. Chỉ load checkpoint từ nguồn tin cậy. Nếu chỉ cần weights cho inference, ưu tiên lưu/load state dict riêng và cân nhắc weights_only=True khi phù hợp.

11. Reproducibility

Reproducibility trong deep learning là mục tiêu thực dụng, không phải lời hứa tuyệt đối. Một số CUDA operation có thể nondeterministic, version library khác nhau cũng có thể làm metric lệch nhẹ.

Tối thiểu nên có:

def seed_everything(seed: int) -> None:
    import random
    import torch

    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

Trade-off:

• Deterministic setting giúp debug và audit.
• Có thể chậm hơn vì tắt benchmark chọn thuật toán nhanh nhất.
• Production training lớn đôi khi ưu tiên throughput, nhưng vẫn phải version data/config/code để so sánh run.

Nên lưu:

• Seed.
• Git commit hoặc code version.
• Package version.
• Dataset snapshot hoặc query/data hash.
• Train/validation/test split.
• Config optimizer/scheduler.
• Metric và threshold.

12. Logging tối thiểu

Log nên giúp trả lời nhanh các câu hỏi:

• Model có học không?
• Có overfit không?
• Learning rate hiện tại là bao nhiêu?
• Gradient có exploding không?
• Epoch này tốn bao lâu?
• Checkpoint nào đang là best?
• Run này dùng config nào?

Log tối thiểu theo epoch:

epoch=7
train_loss=0.3124
val_loss=0.3501
val_f1=0.8421
lr=0.000750
grad_norm=0.91
seconds=12.44
stale_epochs=1

Production nên dùng structured logging hoặc experiment tracking như MLflow, Weights & Biases, TensorBoard hoặc hệ thống nội bộ. Với bài học này, JSON line hoặc print có format ổn định là đủ.

13. Failure modes và cách debug

Debug workflow nên theo thứ tự:

1. Chạy trên subset nhỏ để giảm feedback loop.
2. Overfit một batch.
3. Tắt AMP và scheduler.
4. Log shape/dtype/device ở boundary.
5. Kiểm tra data split và label.
6. Thêm lại từng optimization một: clipping, scheduler, AMP, DataLoader workers.

14. Production readiness

Training loop trong bài dùng được làm nền cho production, nhưng production thực sự cần nhiều lớp xung quanh:

Nếu thiếu các điều kiện này, training loop vẫn có giá trị học tập hoặc prototype, nhưng chưa nên gọi là production training pipeline.

15. Tóm tắt quyết định kỹ thuật

• Bắt đầu bằng PyTorch loop tự viết để hiểu đúng behavior.
• Dùng BCEWithLogitsLoss cho binary classification logits.
• Dùng AdamW làm default cho prototype/fine-tuning hiện đại.
• Dùng ReduceLROnPlateau khi validation loss là tín hiệu chính; dùng warmup/cosine hoặc OneCycleLR khi biết total steps và training dài.
• Dùng gradient clipping khi có loss spike, RNN/Transformer hoặc fine-tuning không ổn định.
• Dùng AMP sau khi FP32 chạy đúng và có benchmark.
• Save best checkpoint theo validation metric.
• Log đủ để debug, reproduce và so sánh run.

Bài tập

Mục tiêu thực hành

Bạn sẽ viết một training job PyTorch cho binary classifier trên synthetic dataset. Dataset synthetic giúp bài chạy offline, không cần download dữ liệu, nhưng training loop được tổ chức theo style gần production:

• Config bằng dataclass.
• Seed và device rõ ràng.
• Dataset/DataLoader riêng.
• Train/eval loop riêng.
• model.train(), model.eval(), torch.inference_mode().
• optimizer.zero_grad(set_to_none=True).
• loss.backward(), gradient clipping, optimizer.step().
• AdamW optimizer.
• ReduceLROnPlateau scheduler.
• Optional mixed precision bằng torch.amp.
• Early stopping.
• Checkpoint best model.
• Metric logging theo JSON line.

Setup

Cài PyTorch theo môi trường của bạn. Nếu đã có PyTorch thì bỏ qua bước này.

pip install torch

Chạy script có sẵn trong folder bài học:

python lessions/day-11-training-loop-optimizer-scheduler/day11_training_loop.py

Nếu có CUDA GPU, script tự dùng GPU. Nếu không có, script chạy CPU.

Script mẫu hoàn chỉnh

Nội dung dưới đây giống file day11_training_loop.py, được đặt trong Markdown để bạn đọc và annotate dễ hơn.

from __future__ import annotations

import json
import random
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import torch
from torch import Tensor, nn
from torch.utils.data import DataLoader, Dataset, random_split


@dataclass(frozen=True)
class TrainConfig:
    seed: int = 42
    n_samples: int = 12_000
    input_dim: int = 8
    hidden_dim: int = 64
    batch_size: int = 256
    epochs: int = 40
    lr: float = 3e-3
    weight_decay: float = 1e-2
    dropout: float = 0.15
    grad_clip_norm: float = 1.0
    early_stopping_patience: int = 5
    min_delta: float = 1e-4
    num_workers: int = 0
    use_amp: bool = True
    threshold: float = 0.5
    checkpoint_path: str = "artifacts/day11_best_checkpoint.pt"


def seed_everything(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def select_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def log_event(event: str, payload: dict[str, Any]) -> None:
    record = {"event": event, **payload}
    print(json.dumps(record, ensure_ascii=False, sort_keys=True))


class SyntheticTicketDataset(Dataset[tuple[Tensor, Tensor]]):
    """Synthetic binary classification data with non-linear signal.

    Features can be read as normalized ticket/review signals:
    length, sentiment cues, refund cue, urgency cue and noise columns.
    """

    def __init__(self, n_samples: int, input_dim: int, seed: int) -> None:
        if input_dim < 6:
            raise ValueError("input_dim must be at least 6")

        generator = torch.Generator().manual_seed(seed)
        features = torch.randn(n_samples, input_dim, generator=generator)

        raw_score = (
            1.4 * features[:, 0]
            - 1.1 * features[:, 1]
            + 0.9 * features[:, 2] * features[:, 3]
            + 0.7 * torch.relu(features[:, 4])
            - 0.6 * features[:, 5].abs()
            - 0.15
        )
        probs = torch.sigmoid(raw_score)
        labels = torch.bernoulli(probs, generator=generator).unsqueeze(1)

        self.features = features.float()
        self.labels = labels.float()

    def __len__(self) -> int:
        return self.features.shape[0]

    def __getitem__(self, index: int) -> tuple[Tensor, Tensor]:
        return self.features[index], self.labels[index]


class BinaryClassifier(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, dropout: float) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(self, x: Tensor) -> Tensor:
        if x.ndim != 2:
            raise ValueError(f"Expected input shape (batch, features), got {tuple(x.shape)}")
        return self.net(x)


def seed_worker(worker_id: int) -> None:
    worker_seed = torch.initial_seed() % 2**32
    random.seed(worker_seed)


def build_loaders(cfg: TrainConfig, device: torch.device) -> tuple[DataLoader, DataLoader, DataLoader]:
    dataset = SyntheticTicketDataset(
        n_samples=cfg.n_samples,
        input_dim=cfg.input_dim,
        seed=cfg.seed,
    )

    n_train = int(0.70 * len(dataset))
    n_val = int(0.15 * len(dataset))
    n_test = len(dataset) - n_train - n_val
    split_generator = torch.Generator().manual_seed(cfg.seed)
    train_ds, val_ds, test_ds = random_split(
        dataset,
        [n_train, n_val, n_test],
        generator=split_generator,
    )

    pin_memory = device.type == "cuda"
    common = {
        "batch_size": cfg.batch_size,
        "num_workers": cfg.num_workers,
        "pin_memory": pin_memory,
        "persistent_workers": cfg.num_workers > 0,
        "worker_init_fn": seed_worker if cfg.num_workers > 0 else None,
    }

    train_loader = DataLoader(
        train_ds,
        shuffle=True,
        generator=torch.Generator().manual_seed(cfg.seed),
        **common,
    )
    val_loader = DataLoader(val_ds, shuffle=False, **common)
    test_loader = DataLoader(test_ds, shuffle=False, **common)
    return train_loader, val_loader, test_loader


def binary_metrics(logits: Tensor, targets: Tensor, threshold: float) -> dict[str, float]:
    probs = torch.sigmoid(logits)
    preds = (probs >= threshold).int()
    y_true = targets.int()

    tp = int(((preds == 1) & (y_true == 1)).sum().item())
    tn = int(((preds == 0) & (y_true == 0)).sum().item())
    fp = int(((preds == 1) & (y_true == 0)).sum().item())
    fn = int(((preds == 0) & (y_true == 1)).sum().item())

    precision = tp / max(tp + fp, 1)
    recall = tp / max(tp + fn, 1)
    f1 = 2 * precision * recall / max(precision + recall, 1e-12)
    accuracy = (tp + tn) / max(tp + tn + fp + fn, 1)

    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }


def current_lr(optimizer: torch.optim.Optimizer) -> float:
    return float(optimizer.param_groups[0]["lr"])


def resolve_output_path(path_text: str) -> Path:
    path = Path(path_text)
    if path.is_absolute():
        return path
    return Path(__file__).resolve().parent / path


def train_one_epoch(
    model: nn.Module,
    loader: DataLoader,
    loss_fn: nn.Module,
    optimizer: torch.optim.Optimizer,
    scaler: torch.amp.GradScaler,
    device: torch.device,
    cfg: TrainConfig,
) -> tuple[float, dict[str, float]]:
    model.train()
    amp_enabled = cfg.use_amp and device.type == "cuda"
    total_loss = 0.0
    total_examples = 0
    grad_norms: list[float] = []
    all_logits: list[Tensor] = []
    all_targets: list[Tensor] = []

    for xb, yb in loader:
        xb = xb.to(device, non_blocking=True)
        yb = yb.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)
        with torch.amp.autocast(
            device_type=device.type,
            dtype=torch.float16,
            enabled=amp_enabled,
        ):
            logits = model(xb)
            loss = loss_fn(logits, yb)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            max_norm=cfg.grad_clip_norm,
        )
        scaler.step(optimizer)
        scaler.update()

        batch_size = xb.shape[0]
        total_loss += float(loss.detach().item()) * batch_size
        total_examples += batch_size
        grad_norms.append(float(grad_norm.detach().cpu().item()))
        all_logits.append(logits.detach().cpu())
        all_targets.append(yb.detach().cpu())

    metrics = binary_metrics(
        torch.cat(all_logits),
        torch.cat(all_targets),
        threshold=cfg.threshold,
    )
    metrics["grad_norm"] = sum(grad_norms) / max(len(grad_norms), 1)
    return total_loss / max(total_examples, 1), metrics


@torch.inference_mode()
def evaluate(
    model: nn.Module,
    loader: DataLoader,
    loss_fn: nn.Module,
    device: torch.device,
    cfg: TrainConfig,
) -> tuple[float, dict[str, float]]:
    model.eval()
    total_loss = 0.0
    total_examples = 0
    all_logits: list[Tensor] = []
    all_targets: list[Tensor] = []

    for xb, yb in loader:
        xb = xb.to(device, non_blocking=True)
        yb = yb.to(device, non_blocking=True)

        logits = model(xb)
        loss = loss_fn(logits, yb)

        batch_size = xb.shape[0]
        total_loss += float(loss.item()) * batch_size
        total_examples += batch_size
        all_logits.append(logits.cpu())
        all_targets.append(yb.cpu())

    metrics = binary_metrics(
        torch.cat(all_logits),
        torch.cat(all_targets),
        threshold=cfg.threshold,
    )
    return total_loss / max(total_examples, 1), metrics


def save_checkpoint(
    path: Path,
    epoch: int,
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.ReduceLROnPlateau,
    scaler: torch.amp.GradScaler,
    cfg: TrainConfig,
    best_val_loss: float,
    val_metrics: dict[str, float],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "scaler_state_dict": scaler.state_dict(),
            "config": asdict(cfg),
            "best_val_loss": best_val_loss,
            "val_metrics": val_metrics,
        },
        path,
    )
    path.with_suffix(".config.json").write_text(
        json.dumps(asdict(cfg), ensure_ascii=False, indent=2, sort_keys=True),
        encoding="utf-8",
    )


def main() -> None:
    cfg = TrainConfig()
    seed_everything(cfg.seed)
    device = select_device()
    checkpoint_path = resolve_output_path(cfg.checkpoint_path)

    train_loader, val_loader, test_loader = build_loaders(cfg, device)
    model = BinaryClassifier(
        input_dim=cfg.input_dim,
        hidden_dim=cfg.hidden_dim,
        dropout=cfg.dropout,
    ).to(device)

    loss_fn = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=cfg.lr,
        weight_decay=cfg.weight_decay,
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=0.5,
        patience=2,
    )

    amp_enabled = cfg.use_amp and device.type == "cuda"
    scaler = torch.amp.GradScaler(device="cuda", enabled=amp_enabled)

    log_event(
        "run_started",
        {
            "device": str(device),
            "amp_enabled": amp_enabled,
            "config": asdict(cfg),
            "train_batches": len(train_loader),
            "val_batches": len(val_loader),
        },
    )

    best_val_loss = float("inf")
    stale_epochs = 0

    for epoch in range(1, cfg.epochs + 1):
        started = time.perf_counter()
        train_loss, train_metrics = train_one_epoch(
            model=model,
            loader=train_loader,
            loss_fn=loss_fn,
            optimizer=optimizer,
            scaler=scaler,
            device=device,
            cfg=cfg,
        )
        val_loss, val_metrics = evaluate(
            model=model,
            loader=val_loader,
            loss_fn=loss_fn,
            device=device,
            cfg=cfg,
        )

        scheduler.step(val_loss)
        improved = val_loss < best_val_loss - cfg.min_delta
        if improved:
            best_val_loss = val_loss
            stale_epochs = 0
            save_checkpoint(
                path=checkpoint_path,
                epoch=epoch,
                model=model,
                optimizer=optimizer,
                scheduler=scheduler,
                scaler=scaler,
                cfg=cfg,
                best_val_loss=best_val_loss,
                val_metrics=val_metrics,
            )
        else:
            stale_epochs += 1

        log_event(
            "epoch_finished",
            {
                "epoch": epoch,
                "train_loss": round(train_loss, 6),
                "val_loss": round(val_loss, 6),
                "train_f1": round(train_metrics["f1"], 6),
                "val_f1": round(val_metrics["f1"], 6),
                "val_precision": round(val_metrics["precision"], 6),
                "val_recall": round(val_metrics["recall"], 6),
                "lr": current_lr(optimizer),
                "grad_norm": round(train_metrics["grad_norm"], 6),
                "seconds": round(time.perf_counter() - started, 3),
                "best_val_loss": round(best_val_loss, 6),
                "stale_epochs": stale_epochs,
                "checkpoint_saved": improved,
            },
        )

        if stale_epochs >= cfg.early_stopping_patience:
            log_event(
                "early_stopping",
                {
                    "epoch": epoch,
                    "best_val_loss": round(best_val_loss, 6),
                    "patience": cfg.early_stopping_patience,
                },
            )
            break

    if not checkpoint_path.exists():
        raise RuntimeError("No checkpoint was saved. Check validation loop and config.")

    # Full training checkpoints should only be loaded from trusted storage.
    checkpoint = torch.load(
        checkpoint_path,
        map_location=device,
        weights_only=False,
    )
    model.load_state_dict(checkpoint["model_state_dict"])
    test_loss, test_metrics = evaluate(
        model=model,
        loader=test_loader,
        loss_fn=loss_fn,
        device=device,
        cfg=cfg,
    )
    log_event(
        "test_finished",
        {
            "checkpoint_epoch": int(checkpoint["epoch"]),
            "test_loss": round(test_loss, 6),
            "test_metrics": {k: round(v, 6) for k, v in test_metrics.items()},
            "best_val_loss": round(float(checkpoint["best_val_loss"]), 6),
        },
    )


if __name__ == "__main__":
    main()

Kết quả kỳ vọng

Bạn sẽ thấy log dạng JSON line:

{"event": "epoch_finished", "epoch": 1, "train_loss": 0.657812, "val_loss": 0.617493, "val_f1": 0.701234, "lr": 0.003, "checkpoint_saved": true}

Metric cụ thể có thể khác theo hardware và version PyTorch, nhưng xu hướng hợp lý là:

• Train loss giảm qua nhiều epoch.
• Validation loss giảm rồi chững lại.
• Scheduler có thể giảm learning rate nếu validation loss plateau.
• Early stopping có thể dừng trước epochs.
• File checkpoint được tạo ở lessions/day-11-training-loop-optimizer-scheduler/artifacts/day11_best_checkpoint.pt nếu chạy từ repo root.

Bài tập 1: Chạy baseline và đọc log

Chạy script với config mặc định.

Ghi lại:

Giải thích ngắn: vì sao không dùng test set để quyết định checkpoint tốt nhất?

Bài tập 2: Thử learning rate

Chạy 3 cấu hình:

Câu hỏi:

1. Run nào converge nhanh nhất?
2. Run nào validation loss dao động mạnh nhất?
3. Nếu loss thành NaN hoặc metric tệ, bạn debug theo thứ tự nào?

Bài tập 3: Thử batch size

Chạy 3 cấu hình:

Câu hỏi:

1. Batch size lớn có luôn tốt hơn không?
2. Nếu tăng batch size làm validation F1 giảm, bạn thử điều chỉnh gì?
3. Trên GPU, bạn quan sát memory và utilization bằng công cụ nào?

Bài tập 4: Thử regularization

Thử các giá trị:

Câu hỏi:

1. Run nào có dấu hiệu overfit?
2. Run nào có dấu hiệu underfit?
3. Vì sao dropout quá cao có thể làm model học chậm?

Bài tập 5: Tắt từng optimization để debug

Thử lần lượt:

1. use_amp=False.
2. grad_clip_norm=10.0.
3. early_stopping_patience=20.
4. Scheduler patience=5 thay vì 2.

Ghi lại ảnh hưởng đến:

• Stability.
• Thời gian chạy.
• Validation F1.
• Số epoch trước khi dừng.

Mục tiêu không phải tìm con số đẹp nhất. Mục tiêu là hiểu mỗi cơ chế ảnh hưởng thế nào đến training behavior.

Bài tập 6: Production review

Trả lời các câu hỏi sau như một review trước khi đưa training job vào production:

1. Dataset thật sẽ được version bằng gì?
2. Split train/validation/test cần theo random, user hay time? Vì sao?
3. Metric business chính là precision, recall, F1 hay PR-AUC?
4. Threshold có nên giữ 0.5 không?
5. Checkpoint cần lưu ở local disk, object storage hay model registry?
6. Log hiện tại có đủ để debug loss spike không?
7. Có rủi ro log PII không?
8. Nếu job bị kill giữa chừng, cần resume từ đâu?
9. Nếu model mới tệ hơn production model hiện tại, rollback thế nào?
10. Dùng được trong production không? Nếu có thì cần thêm điều kiện gì?

Gợi ý lời giải production readiness

Training job này có thể làm nền cho production vì đã có seed, config, train/validation/test split, checkpoint best model, metric logging, scheduler, gradient clipping và early stopping.

Nhưng để dùng production thật, cần bổ sung:

• Dataset version và schema validation.
• Data quality check, leakage check và class distribution report.
• Experiment tracking hoặc structured logs được thu thập tập trung.
• Artifact storage đáng tin cậy thay vì chỉ local artifacts/.
• Unit test cho dataset, metric, checkpoint load và shape contract.
• Threshold tuning theo validation set và business cost.
• Package lockfile, Docker image hoặc environment reproducible.
• Alert khi loss NaN, validation metric tụt mạnh hoặc training cost vượt budget.
• Quy trình approval trước khi promote model sang inference service.

Câu hỏi tự kiểm tra

1. Vì sao optimizer.zero_grad(set_to_none=True) nên gọi trước batch mới?
2. model.eval() khác gì với torch.no_grad()?
3. Với ReduceLROnPlateau, vì sao cần gọi scheduler.step(val_loss) thay vì scheduler.step()?
4. Nếu dùng AMP và gradient clipping, vì sao cần scaler.unscale_(optimizer) trước clip_grad_norm_?
5. Checkpoint để resume training khác gì checkpoint chỉ để inference?
6. Vì sao không nên chọn model theo test metric?
7. Khi validation loss tăng nhưng train loss giảm, bạn nghĩ đến những nguyên nhân nào?
8. Khi GPU utilization thấp, bạn kiểm tra DataLoader như thế nào?

Checklist hoàn thành exercise

• Chạy được script trên CPU hoặc CUDA.
• Có log run_started, epoch_finished, early_stopping nếu xảy ra và test_finished.
• Có checkpoint tại lessions/day-11-training-loop-optimizer-scheduler/artifacts/day11_best_checkpoint.pt nếu chạy từ repo root.
• Thay đổi learning rate và giải thích được khác biệt.
• Thay đổi batch size và giải thích được trade-off performance.
• Thay đổi dropout/weight decay và nhận diện overfit/underfit.
• Tắt AMP để debug và hiểu khi nào nên bật lại.
• Viết được production review ngắn cho training job.