Bài 41: Experiment Tracking với Weights & Biases (W&B)

Sau bài này bạn sẽ:

• ✅ Hiểu W&B là gì, khác gì MLflow về mô hình hosting và license
• ✅ Cài đặt, login, và chạy experiment tracking đầu tiên với wandb 0.17+
• ✅ Log metrics, images, tables, và artifacts đúng cách
• ✅ Cấu hình và chạy Sweeps để tự động tìm hyperparameter tốt
• ✅ Tích hợp W&B với PyTorch Lightning, Hugging Face Transformers, Keras
• ✅ Biết khi nào dùng W&B, khi nào dùng MLflow

Weights & Biases (W&B) là platform experiment tracking và collaboration cho machine learning. Điểm khác biệt chính so với MLflow:

• Hosted by default: dữ liệu lưu trên cloud W&B (US/EU), không cần tự vận hành server. MLflow yêu cầu tự host nếu cần persistence.
• Proprietary, nhưng free tier rộng: cá nhân và project open source được dùng miễn phí không giới hạn run. Team collaboration cần paid plan.
• Self-host có sẵn qua W&B Server (enterprise license) — cài on-prem nếu cần data locality.

W&B phù hợp khi:

• Không muốn ops server tracking riêng.
• Cần dashboard đẹp, dễ share với stakeholder.
• Team research cần collaboration và comment trực tiếp trên chart.
• Dùng PyTorch Lightning hoặc Hugging Face Transformers nặng — W&B có integration sâu.
• Cần hyperparameter sweep mạnh (Bayesian) tích hợp native.

pip install wandb          # wandb 0.17+ tính đến 2025
wandb login               # dán API key khi được hỏi

API key lấy từ https://wandb.ai/authorize sau khi tạo account. Key được lưu vào ~/.netrc — không cần login lại ở các lần sau.

Cách khác dùng biến môi trường, phù hợp cho CI/CD hoặc container:

export WANDB_API_KEY=your_key_here

Kiểm tra version:

python -c "import wandb; print(wandb.__version__)"
# 0.17.x

Lưu ý bảo mật: không hard-code API key trong source code commit lên git. Dùng env var hoặc secret manager.

import wandb
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

wandb.init(
    project="iris-classification",
    config={"C": 1.0, "max_iter": 1000},
)
config = wandb.config

model = LogisticRegression(C=config.C, max_iter=config.max_iter)
model.fit(X_train, y_train)
acc = accuracy_score(y_test, model.predict(X_test))

wandb.log({"accuracy": acc})
wandb.finish()

Sau khi chạy, terminal in ra link trực tiếp đến run trên W&B UI dạng https://wandb.ai/<username>/iris-classification/runs/<run_id>.

Trong UI có thể xem params từ config, metric accuracy, system metrics (CPU/GPU usage, RAM), và git commit nếu chạy trong git repo.

run = wandb.init(
    project="my-project",        # tên project trên W&B cloud
    name="experiment-v3",        # tên run (random nếu không set)
    config={                     # hyperparameter — accessible qua wandb.config
        "lr": 1e-3,
        "batch_size": 32,
        "epochs": 20,
    },
    tags=["baseline", "resnet"],  # list tag để filter trên UI
    group="cross-val-fold",       # gom nhiều run vào cùng 1 group (vd cross-validation)
    job_type="train",             # "train" | "eval" | "preprocess" | ...
    mode="online",                # "online" (default) | "offline" | "disabled"
)

mode — 3 Chế Độ

• online (default): log real-time lên W&B cloud. Cần kết nối mạng liên tục.
• offline: lưu local vào ./wandb/. Sync lên cloud sau bằng wandb sync ./wandb/<run_dir>. Phù hợp khi train ở nơi mạng không ổn định.
• disabled: tắt hoàn toàn W&B — mọi lệnh wandb.* trở thành no-op. Hữu ích khi test code mà không muốn tạo run.

config vs wandb.config

config truyền vào init() là dict khởi tạo. wandb.config là object có thể đọc lại (dot notation) và cũng được sweep agent ghi đè khi chạy sweep:

wandb.init(project="test", config={"lr": 1e-3})
config = wandb.config
print(config.lr)   # 0.001

Khi dùng sweep, agent tự điền giá trị vào wandb.config trước khi code chạy — không cần parse argument thủ công.

Metrics Theo Epoch

for epoch in range(epochs):
    train_loss = train_one_epoch(model, loader, optimizer)
    val_loss = validate(model, val_loader)

    wandb.log({
        "train_loss": train_loss,
        "val_loss": val_loss,
        "epoch": epoch,
    })
    # step tự tăng mỗi lần gọi wandb.log()

Override step thủ công nếu cần đồng bộ với epoch number cụ thể:

wandb.log({"val_loss": val_loss}, step=epoch)

Lưu ý: step phải tăng đơn điệu (monotonically increasing). Log step không tăng làm chart trên UI hiển thị sai thứ tự điểm dữ liệu.

Log Image

import numpy as np
from PIL import Image

# Từ numpy array
img_array = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
wandb.log({"sample_image": wandb.Image(img_array, caption="generated sample")})

# Từ PIL Image
pil_img = Image.open("test.jpg")
wandb.log({"input": wandb.Image(pil_img)})

Log Table

import pandas as pd

df = pd.DataFrame({
    "image_path": test_paths,
    "true_label": y_true,
    "predicted": y_pred,
    "confidence": confidences,
})
wandb.log({"predictions": wandb.Table(dataframe=df)})

W&B Table render trong UI thành bảng tương tác — có thể filter, sort, và tạo chart trực tiếp từ dữ liệu. Phù hợp để phân tích error case.

Các Media Type Khác

• wandb.Audio(audio_array, sample_rate=44100)
• wandb.Video(frames_array)
• wandb.Histogram(values) — phân phối giá trị
• wandb.plot.confusion_matrix(y_true, y_pred, class_names)

Artifact là cơ chế versioning file trong W&B: model checkpoint, dataset, processed output. Mỗi lần save tạo version mới tự động (v0, v1, ...). Tag latest trỏ vào version mới nhất.

Lưu Artifact

import wandb
import pickle

# Ví dụ: lưu model sau train
with open("model.pkl", "wb") as f:
    pickle.dump(model, f)

artifact = wandb.Artifact(
    name="iris-logistic-model",   # tên artifact
    type="model",                  # "model" | "dataset" | "result"
    description="LogisticRegression C=1.0 trained on Iris",
)
artifact.add_file("./model.pkl")
wandb.log_artifact(artifact)
# W&B tự gán version: iris-logistic-model:v0, v1, v2, ...

Thêm thư mục:

artifact.add_dir("./checkpoints/")

Load Artifact Trong Run Khác

wandb.init(project="iris-classification", job_type="eval")

# Lấy version mới nhất
artifact = wandb.use_artifact("iris-logistic-model:latest")
artifact_dir = artifact.download()   # tải về thư mục local tạm

# Hoặc lấy version cụ thể
artifact_v1 = wandb.use_artifact("iris-logistic-model:v1")
artifact_dir_v1 = artifact_v1.download()

wandb.use_artifact() tạo liên kết lineage: run hiện tại "consume" artifact đó. W&B UI hiển thị graph upstream/downstream artifact — biết model nào dùng dataset nào, train run nào sinh ra artifact nào.

Artifact Dataset

artifact = wandb.Artifact("iris-dataset", type="dataset", metadata={"source": "sklearn"})
artifact.add_dir("./data/")
wandb.log_artifact(artifact)

Sweeps là cơ chế hyperparameter search tích hợp native trong W&B. Controller (W&B cloud) quản lý search strategy và phân phối config cho từng agent. Có thể chạy nhiều agent song song trên nhiều máy.

Bước 1: Tạo sweep.yaml

# sweep.yaml
method: bayes           # "grid" | "random" | "bayes"
metric:
  name: val_loss
  goal: minimize
parameters:
  lr:
    distribution: log_uniform_values
    min: 1e-5
    max: 1e-2
  batch_size:
    values: [16, 32, 64, 128]
  dropout:
    distribution: uniform
    min: 0.0
    max: 0.5

Ba phương pháp tìm kiếm:

• grid: thử toàn bộ tổ hợp. Thực tế chỉ dùng khi số params nhỏ.
• random: sample ngẫu nhiên. Tốt khi không biết vùng tốt.
• bayes: dùng Bayesian optimization — ưu tiên vùng config có khả năng tốt hơn dựa trên kết quả run trước. Hiệu quả hơn random khi budget run có hạn.

Bước 2: Khởi Tạo Sweep

wandb sweep sweep.yaml
# Output:
# wandb: Creating sweep with ID: abc123xyz
# wandb: Run sweep agent with: wandb agent username/project/abc123xyz

Bước 3: Viết Training Function

import wandb
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split, cross_val_score

def train():
    wandb.init()               # agent tự điền config vào wandb.config
    config = wandb.config

    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    model = LogisticRegression(
        C=1.0 / config.lr,      # C là inverse của regularization strength
        max_iter=1000,
    )
    model.fit(X_train, y_train)

    val_loss = 1.0 - model.score(X_test, y_test)
    wandb.log({"val_loss": val_loss})


if __name__ == "__main__":
    train()

Bước 4: Chạy Agent

wandb agent username/project/abc123xyz
# Chạy nhiều agent song song trên máy khác (cùng sweep ID):
# wandb agent username/project/abc123xyz  # terminal khác hoặc máy khác

Mỗi agent lấy một config từ controller, chạy train(), log kết quả, rồi lấy config tiếp. Agent chạy cho đến khi hết run budget hoặc bị dừng tay.

Hoặc Dùng Python API

sweep_config = {
    "method": "bayes",
    "metric": {"name": "val_loss", "goal": "minimize"},
    "parameters": {
        "lr": {"distribution": "log_uniform_values", "min": 1e-5, "max": 1e-2},
        "batch_size": {"values": [16, 32, 64]},
    },
}

sweep_id = wandb.sweep(sweep_config, project="iris-sweep")
wandb.agent(sweep_id, function=train, count=30)  # chạy tối đa 30 run

PyTorch Lightning

from lightning.pytorch.loggers import WandbLogger
from lightning import Trainer

logger = WandbLogger(
    project="my-project",
    name="lightning-run-v1",
    log_model=True,          # tự log model artifact cuối epoch tốt nhất
)

trainer = Trainer(
    max_epochs=20,
    logger=logger,
)
trainer.fit(model, datamodule=dm)

WandbLogger tự log train/val loss mỗi step, learning rate schedule, system metrics (GPU util, VRAM), và checkpoint artifact nếu log_model=True.

Hugging Face Transformers

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    report_to="wandb",         # bật W&B integration
    run_name="bert-finetune",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
)
trainer.train()
# W&B tự log train/eval loss, learning rate theo step

Tắt W&B logging trong Transformers nếu không cần:

export WANDB_DISABLED=true

Keras

import wandb
from wandb.keras import WandbCallback

wandb.init(project="keras-mnist", config={"epochs": 10, "batch_size": 64})

model.fit(
    X_train, y_train,
    epochs=wandb.config.epochs,
    batch_size=wandb.config.batch_size,
    validation_split=0.1,
    callbacks=[WandbCallback()],   # tự log metrics mỗi epoch + model weights histogram
)
wandb.finish()

Reports là tính năng không có trong MLflow. Một Report là tài liệu gồm markdown + chart + bảng số liệu lấy trực tiếp từ runs — có thể share link cho người không có W&B account.

Tạo Report từ UI:

1. Vào project → tab "Reports" → "New report".
2. Viết markdown, chèn chart từ runs (drag từ panel runs).
3. Pin chart từ run cụ thể, embed code snippet.
4. Publish → lấy link share (public hoặc team-only).

Tạo Report từ Python API (wandb-workspaces, thư viện riêng):

import wandb_workspaces.reports.v2 as wr

report = wr.Report(
    project="iris-classification",
    title="Iris Experiment Weekly Summary",
    description="So sánh LR=1e-3 vs 1e-4 trên 3 kiến trúc",
)

report.blocks = [
    wr.H1(text="Kết Quả"),
    wr.RunsetPanel(
        runsets=[
            wr.Runset(project="iris-classification", filters={"tags": {"$in": ["v2"]}})
        ]
    ),
]

report.save()
print(report.url)

Use case thực tế: weekly experiment review — team lead tạo report tóm tắt tuần, link vào Slack/Notion thay vì screenshot.

W&B Table là cách log dataset + prediction để phân tích lỗi trực tiếp trong UI mà không cần export CSV rồi mở notebook riêng.

import wandb
import numpy as np

wandb.init(project="error-analysis")

# Log prediction table
columns = ["image", "true_label", "predicted", "confidence"]
table_data = []

for img, true_label, pred, conf in zip(test_images, y_true, y_pred, confidences):
    table_data.append([
        wandb.Image(img),
        true_label,
        pred,
        round(float(conf), 4),
    ])

table = wandb.Table(columns=columns, data=table_data)
wandb.log({"prediction_table": table})

wandb.finish()

Trên W&B UI, Table có thể:

• Filter theo điều kiện (vd predicted != true_label để xem error case).
• Sort theo confidence để xem prediction chắc chắn nhất / kém nhất.
• Tạo histogram distribution từ cột bất kỳ.
• Group by label để đếm phân phối.

Đây là điểm cần xem xét kỹ trước khi dùng W&B ở môi trường production:

• Default: metrics, params, artifacts lên W&B cloud (US/EU). Dữ liệu thuộc account của bạn, W&B có cam kết không dùng data của user.
• W&B Server (self-host): cài on-prem qua Kubernetes hoặc Docker. Data ở lại trong hạ tầng của tổ chức. Cần enterprise license.
• Anonymous mode: log không cần đăng nhập — W&B tạo anonymous run, claim sau 7 ngày. Dùng cho demo nhanh:

wandb.init(project="demo", anonymous="allow")

Cần chú ý khi log data nhạy cảm:

• Không log raw text chứa PII (tên, email, CCCD, ...) lên cloud.
• Không log dataset y tế hoặc tài chính khi chưa có clearance từ DPO/legal.
• Artifact chứa model weight là an toàn trong hầu hết trường hợp — nhưng kiểm tra xem model có memorize training data không (membership inference attack).
• Cân nhắc dùng mode="offline" khi train trên data nhạy cảm, rồi chỉ sync metrics (không sync artifact dataset) sau.

Chọn MLflow khi

• Open source là yêu cầu cứng (compliance, procurement).
• Cần giữ toàn bộ data on-prem không gửi ra ngoài.
• Đã có infra Databricks hoặc Azure ML — MLflow là default tracking ở đó.
• Cần autologging nhanh không cần viết thêm code tích hợp.
• Cần Model Registry chuẩn để quản lý stage/version — bài 42 sẽ đào sâu phần này.

Chọn W&B khi

• Team nhỏ hoặc vừa, không có người ops MLflow server.
• Research team cần chia sẻ kết quả qua Report link thay vì screenshot.
• Project dùng PyTorch Lightning hoặc HF Transformers nhiều — integration W&B mượt hơn.
• Cần Bayesian sweep tích hợp không cần cài thêm Optuna/Ray Tune.
• Collaboration real-time: nhiều người cùng xem dashboard, comment trên chart.

Dùng Cả Hai

Không hiếm trường hợp dùng song song: W&B cho experiment tracking hàng ngày (UI đẹp, dễ share), MLflow Model Registry cho quản lý model production (on-prem, integrate với serving pipeline). Hai tool không xung đột — cùng một training run có thể log cả hai.

1. Quên wandb.finish()

# SAI — run ở trạng thái "running" mãi trong UI nếu script crash hoặc thoát không sạch
wandb.init(project="test")
wandb.log({"acc": 0.9})
# Script kết thúc mà không gọi finish()

# ĐÚNG — dùng context manager hoặc gọi finish() cuối script
with wandb.init(project="test") as run:
    wandb.log({"acc": 0.9})
# finish() tự gọi khi ra khỏi with block

2. Log Step Không Đơn Điệu

# SAI — step giảm làm chart không đúng thứ tự
wandb.log({"loss": 0.5}, step=10)
wandb.log({"loss": 0.3}, step=5)   # step nhỏ hơn step trước

# ĐÚNG — step chỉ tăng
for i, loss in enumerate(losses):
    wandb.log({"loss": loss}, step=i)

3. Log File Quá Lớn Mỗi Step

Log image 10MB mỗi step, 1000 step = 10GB artifact quota. Giảm tần suất log media:

for step, (img, pred) in enumerate(eval_results):
    if step % 50 == 0:    # chỉ log mỗi 50 step
        wandb.log({"sample": wandb.Image(img)}, step=step)

4. API Key Trong Code

# KHÔNG BAO GIỜ làm thế này
wandb.login(key="your_secret_api_key_here")  # key sẽ vào git history

# ĐÚNG — dùng env var
# export WANDB_API_KEY=your_key
wandb.login()   # đọc từ env var hoặc ~/.netrc

5. Mạng Yếu Trong mode="online"

Khi mạng không ổn định, wandb.log() block chờ upload. Training bị chậm hoặc crash. Dùng mode="offline" khi train:

# Train với offline mode
export WANDB_MODE=offline
python train.py

# Sau khi train xong, sync lên cloud
wandb sync wandb/offline-run-*/

6. Sweep Không Hội Tụ Vì Run Quá Ít

Bayesian optimization cần đủ data điểm để ước lượng surrogate function. Với không gian 3-4 hyperparameter, cần tối thiểu 20-30 run để Bayesian sweep hoạt động tốt hơn random. Nếu budget ít hơn, dùng method: random thay thế.

Bài 42: Model Registry — versioning model artifact — cách tổ chức, đánh version và quản lý lifecycle (staging, production, archived) cho model artifact.