Transfer Learning

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import DataLoader

NUM_CLASSES = 10  # bài toán mới

# ──────────────────────────────────────────────────────────────
# 1. Feature extraction — đóng băng toàn bộ backbone
# ──────────────────────────────────────────────────────────────
def build_feature_extractor() -> nn.Module:
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
    for param in model.parameters():
        param.requires_grad = False          # đóng băng
    # Thay classifier head cho bài toán mới
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, NUM_CLASSES)
    return model

# ──────────────────────────────────────────────────────────────
# 2. Fine-tune 3 lớp cuối — mở khoá layer4 + fc
# ──────────────────────────────────────────────────────────────
def build_last3_finetune() -> nn.Module:
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
    for param in model.parameters():
        param.requires_grad = False
    for param in model.layer4.parameters():
        param.requires_grad = True
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, NUM_CLASSES)  # trainable theo mặc định
    return model

# ──────────────────────────────────────────────────────────────
# 3. Full fine-tune — mở khoá toàn bộ
# ──────────────────────────────────────────────────────────────
def build_full_finetune() -> nn.Module:
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, NUM_CLASSES)
    # Mặc định tất cả requires_grad=True
    return model

# ──────────────────────────────────────────────────────────────
# 4. Optimizer với LR theo nhóm tham số
#    - Lớp mới (head): LR lớn hơn
#    - Backbone pretrained: LR rất nhỏ
# ──────────────────────────────────────────────────────────────
def make_optimizer(model: nn.Module, strategy: str):
    if strategy == "feature":
        # chỉ head được cập nhật
        return torch.optim.AdamW(model.fc.parameters(), lr=1e-3, weight_decay=1e-4)
    if strategy == "last3":
        return torch.optim.AdamW(
            [
                {"params": model.layer4.parameters(), "lr": 1e-5},
                {"params": model.fc.parameters(), "lr": 1e-3},
            ],
            weight_decay=1e-4,
        )
    # full
    return torch.optim.AdamW(
        [
            {"params": [p for n, p in model.named_parameters()
                        if not n.startswith("fc")], "lr": 1e-5},
            {"params": model.fc.parameters(), "lr": 1e-3},
        ],
        weight_decay=1e-4,
    )

# ──────────────────────────────────────────────────────────────
# 5. Vòng lặp huấn luyện rút gọn
# ──────────────────────────────────────────────────────────────
def train(model, loader: DataLoader, optim, epochs: int = 10, device="cuda"):
    criterion = nn.CrossEntropyLoss()
    model.to(device).train()
    for epoch in range(epochs):
        running = 0.0
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            optim.zero_grad(set_to_none=True)
            logits = model(x)
            loss = criterion(logits, y)
            loss.backward()
            optim.step()
            running += loss.item() * x.size(0)
        print(f"epoch {epoch+1}: loss={running / len(loader.dataset):.4f}")
    return model

import torch
from torch.optim import AdamW

def discriminative_param_groups(model, base_lr: float = 1e-3, decay: float = 0.3):
    """
    Chia ResNet thành 4 nhóm theo độ sâu và gán LR giảm dần vào lớp nông.
    decay = 0.3 nghĩa là mỗi nhóm sâu hơn có LR gấp 1/0.3 ≈ 3.3 lần lớp trước.
    """
    groups = [
        {"params": list(model.conv1.parameters())
                   + list(model.bn1.parameters())
                   + list(model.layer1.parameters()),
         "lr": base_lr * decay ** 3},
        {"params": list(model.layer2.parameters()),
         "lr": base_lr * decay ** 2},
        {"params": list(model.layer3.parameters()),
         "lr": base_lr * decay},
        {"params": list(model.layer4.parameters())
                   + list(model.fc.parameters()),
         "lr": base_lr},
    ]
    return groups


optim = AdamW(discriminative_param_groups(model, base_lr=1e-3),
              weight_decay=1e-4)

# Kết hợp với cosine schedule và warmup
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optim,
    max_lr=[g["lr"] for g in optim.param_groups],
    total_steps=len(train_loader) * EPOCHS,
    pct_start=0.1,
    div_factor=10,
)