Autoencoder

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class Autoencoder(nn.Module):
    """MNIST autoencoder: 784 → 128 → 32 → 8 → 32 → 128 → 784."""

    def __init__(self, latent_dim: int = 8):
        super().__init__()
        self.latent_dim = latent_dim

        # ─── Encoder ───
        self.encoder = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 32),
            nn.ReLU(),
            nn.Linear(32, latent_dim),  # ← Bottleneck!
        )

        # ─── Decoder (đối xứng) ───
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 32),
            nn.ReLU(),
            nn.Linear(32, 128),
            nn.ReLU(),
            nn.Linear(128, 784),
            nn.Sigmoid(),  # Output ∈ [0, 1] — phù hợp với ảnh normalized
        )

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # x: [B, 784] — flatten ảnh 28×28
        z = self.encoder(x)      # [B, latent_dim]
        x_hat = self.decoder(z)  # [B, 784]
        return x_hat, z

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Chỉ lấy latent code — dùng cho giảm chiều / anomaly detection."""
        return self.encoder(x)

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """Generate từ latent — lưu ý: vanilla AE không generative tốt."""
        return self.decoder(z)


# ─── Training loop ───
def train(latent_dim: int = 8, epochs: int = 20, batch_size: int = 256):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.view(-1)),  # flatten 28×28 → 784
    ])
    train_ds = datasets.MNIST("./data", train=True, download=True, transform=transform)
    loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)

    model = Autoencoder(latent_dim=latent_dim).to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    # BCE phù hợp với ảnh normalized [0, 1]
    loss_fn = nn.BCELoss()

    for epoch in range(epochs):
        total = 0.0
        for x, _ in loader:  # bỏ qua nhãn — self-supervised!
            x = x.to(device)
            x_hat, _ = model(x)

            loss = loss_fn(x_hat, x)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total += loss.item() * x.size(0)

        avg = total / len(train_ds)
        print(f"Epoch {epoch + 1:2d}/{epochs} | loss = {avg:.4f}")

    return model


# ─── Sử dụng latent code cho anomaly detection ───
@torch.no_grad()
def anomaly_score(model: Autoencoder, x: torch.Tensor) -> torch.Tensor:
    """Reconstruction error cao → anomaly."""
    x_hat, _ = model(x)
    # Per-sample MSE
    return ((x - x_hat) ** 2).mean(dim=-1)


if __name__ == "__main__":
    model = train(latent_dim=8, epochs=10)

    # Kiểm tra: sample 1 mẫu, xem loss
    x_sample = next(iter(DataLoader(
        datasets.MNIST(
            "./data", train=False, download=True,
            transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Lambda(lambda x: x.view(-1)),
            ]),
        ),
        batch_size=1,
    )))[0]
    scores = anomaly_score(model, x_sample)
    print(f"Anomaly score: {scores.item():.4f}")

"""Denoising Autoencoder — variant quan trọng:
   input = ảnh + nhiễu, target = ảnh sạch.
   Mạng phải hiểu bản chất dữ liệu để lọc nhiễu."""

import torch
import torch.nn as nn


class DenoisingAE(nn.Module):
    """Convolutional Denoising AE cho MNIST."""

    def __init__(self, latent_dim: int = 32):
        super().__init__()

        # Encoder: 28×28 → 14×14 → 7×7 → latent
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1),  # 28 → 14
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),  # 14 → 7
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, latent_dim),
        )

        # Decoder: latent → 7×7 → 14×14 → 28×28
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 32 * 7 * 7),
            nn.ReLU(),
            nn.Unflatten(1, (32, 7, 7)),
            nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 1, 3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid(),
        )

    def forward(self, x_noisy):
        z = self.encoder(x_noisy)
        x_hat = self.decoder(z)
        return x_hat


def add_gaussian_noise(x: torch.Tensor, sigma: float = 0.3) -> torch.Tensor:
    """Thêm nhiễu Gaussian. sigma lớn → nhiễu mạnh."""
    return (x + sigma * torch.randn_like(x)).clamp(0, 1)


def train_denoising(model, loader, optimizer, loss_fn, epochs=20, sigma=0.3):
    model.train()
    for epoch in range(epochs):
        total = 0.0
        for x, _ in loader:
            x_clean = x  # target là ảnh sạch
            x_noisy = add_gaussian_noise(x_clean, sigma=sigma)

            x_hat = model(x_noisy)
            # Loss giữa output và ảnh SẠCH (không phải noisy input!)
            loss = loss_fn(x_hat, x_clean)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total += loss.item() * x.size(0)

        print(f"Epoch {epoch + 1:2d} | loss = {total / len(loader.dataset):.4f}")


# Điểm mấu chốt: mạng không thể chỉ copy input sang output
# (vì input ≠ target). Nó BUỘC phải học lọc nhiễu → biểu diễn robust.
#
# Denoising AE là tiền thân của:
#   - BERT: mask 15% token, dự đoán lại → "denoising" trong NLP
#   - MAE (Masked AE): mask 75% patch ảnh → pretrain cho ViT
#   - Diffusion models: denoising lặp lại thành generative mạnh nhất hiện nay