Regularization

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms

# ───────────────────────────────────────────────────────────────
# 1) L2 weight decay: thêm vào optimizer
#    (Adam + weight decay đúng cách = AdamW — decoupled WD)
# ───────────────────────────────────────────────────────────────
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-4,  # λ cho L2
)

# ───────────────────────────────────────────────────────────────
# 2) L1 regularization: thêm thủ công vào loss
# ───────────────────────────────────────────────────────────────
def l1_penalty(model, lambda_l1=1e-5):
    return lambda_l1 * sum(
        p.abs().sum() for p in model.parameters() if p.requires_grad
    )

# training loop:
loss = criterion(output, target) + l1_penalty(model)

# ───────────────────────────────────────────────────────────────
# 3) ElasticNet: kết hợp L1 + L2
# ───────────────────────────────────────────────────────────────
def elastic_penalty(model, l1=1e-5, l2=1e-4, alpha=0.5):
    l1_term = sum(p.abs().sum() for p in model.parameters())
    l2_term = sum((p ** 2).sum() for p in model.parameters())
    return alpha * l1 * l1_term + (1 - alpha) * l2 * l2_term

# ───────────────────────────────────────────────────────────────
# 4) Dropout + BatchNorm trong model definition
# ───────────────────────────────────────────────────────────────
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.BatchNorm1d(256),   # BN trước activation
    nn.ReLU(),
    nn.Dropout(0.3),       # tắt 30% sau activation
    nn.Linear(256, 128),
    nn.BatchNorm1d(128),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(128, 10),
)

# Quan trọng: bật/tắt Dropout và BN khi train/eval
model.train()  # Dropout ON, BN dùng batch stats
model.eval()   # Dropout OFF, BN dùng running stats

# ───────────────────────────────────────────────────────────────
# 5) Data Augmentation (torchvision)
# ───────────────────────────────────────────────────────────────
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.4, 0.4, 0.4),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
    transforms.RandomErasing(p=0.25),  # Cutout
])

# Mixup (trong training loop)
def mixup_data(x, y, alpha=0.2):
    lam = torch.distributions.Beta(alpha, alpha).sample().item()
    idx = torch.randperm(x.size(0))
    mixed_x = lam * x + (1 - lam) * x[idx]
    y_a, y_b = y, y[idx]
    return mixed_x, y_a, y_b, lam

# ───────────────────────────────────────────────────────────────
# 6) Early Stopping helper
# ───────────────────────────────────────────────────────────────
class EarlyStopping:
    def __init__(self, patience=7, min_delta=0.0, restore_best=True):
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best = restore_best
        self.counter = 0
        self.best_loss = float("inf")
        self.best_state = None
        self.should_stop = False

    def step(self, val_loss, model):
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.counter = 0
            if self.restore_best:
                self.best_state = {
                    k: v.detach().clone()
                    for k, v in model.state_dict().items()
                }
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.should_stop = True
                if self.restore_best and self.best_state is not None:
                    model.load_state_dict(self.best_state)

# Usage:
stopper = EarlyStopping(patience=5)
for epoch in range(100):
    train_one_epoch(model, train_loader, optimizer)
    val_loss = evaluate(model, val_loader)
    stopper.step(val_loss, model)
    if stopper.should_stop:
        print(f"Early stopping at epoch {epoch}")
        break

# ───────────────────────────────────────────────────────────────
# 7) Label Smoothing
# ───────────────────────────────────────────────────────────────
# Cách 1 — có sẵn trong PyTorch 1.10+
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# Cách 2 — tự viết (dùng được cho custom loss)
def smooth_cross_entropy(logits, target, smoothing=0.1):
    n_classes = logits.size(-1)
    log_probs = F.log_softmax(logits, dim=-1)
    # target one-hot
    with torch.no_grad():
        true_dist = torch.zeros_like(log_probs)
        true_dist.fill_(smoothing / (n_classes - 1))
        true_dist.scatter_(1, target.unsqueeze(1), 1 - smoothing)
    return -(true_dist * log_probs).sum(dim=-1).mean()

# ───────────────────────────────────────────────────────────────
# 8) DropConnect (không có trong PyTorch — tự implement)
# ───────────────────────────────────────────────────────────────
class DropConnectLinear(nn.Module):
    def __init__(self, in_features, out_features, p=0.3):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.p = p

    def forward(self, x):
        if self.training and self.p > 0:
            mask = torch.bernoulli(
                torch.full_like(self.linear.weight, 1 - self.p)
            )
            W = self.linear.weight * mask / (1 - self.p)  # inverted
            return F.linear(x, W, self.linear.bias)
        return self.linear(x)