reproduction-bundle/src/train_public_mnist_classifiers.py

#!/usr/bin/env python3
"""Train full-data auxiliary classifiers for the pinned public MNIST DDPM.

These classifiers are intentionally separate from the paper's CMNIST models.
They make it possible to exercise LOGDIFF's exact mutually-exclusive rules on a
publicly downloadable, trained diffusion model rather than on a finite-state
proxy.  Every epoch uses all 60,000 MNIST training examples.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import random
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from diffusers import DDPMScheduler
from torch import nn
from torchvision import datasets, transforms


ROOT = Path(__file__).resolve().parents[2]


class ImageEncoder(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.GroupNorm(8, 32),
            nn.SiLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.GroupNorm(8, 128),
            nn.SiLU(),
            nn.Conv2d(128, 128, 3, stride=2, padding=1),
            nn.GroupNorm(8, 128),
            nn.SiLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def sinusoidal_time(t: torch.Tensor, channels: int = 128) -> torch.Tensor:
    half = channels // 2
    frequencies = torch.exp(
        -math.log(10000.0) * torch.arange(half, device=t.device) / (half - 1)
    )
    angles = t.float().unsqueeze(1) * frequencies.unsqueeze(0)
    return torch.cat([angles.sin(), angles.cos()], dim=1)


class NoisyDigitClassifier(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.encoder = ImageEncoder()
        self.time_mlp = nn.Sequential(nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 128))
        self.head = nn.Linear(256, 10)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        return self.head(torch.cat([self.encoder(x), self.time_mlp(sinusoidal_time(t))], dim=1))


class CleanDigitClassifier(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.encoder = ImageEncoder()
        self.head = nn.Linear(128, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.head(self.encoder(x))


def select_device(name: str) -> torch.device:
    if name != "auto":
        return torch.device(name)
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def accuracy(logits: torch.Tensor, target: torch.Tensor) -> tuple[int, int]:
    return int((logits.argmax(1) == target).sum().item()), target.numel()


@torch.inference_mode()
def evaluate(
    noisy_model: NoisyDigitClassifier,
    clean_model: CleanDigitClassifier,
    loader: torch.utils.data.DataLoader,
    alpha: torch.Tensor,
    device: torch.device,
    seed: int,
) -> dict[str, float]:
    noisy_model.eval()
    clean_model.eval()
    noisy_correct = clean_correct = total = 0
    noisy_loss = clean_loss = 0.0
    generator = torch.Generator(device="cpu").manual_seed(seed)
    for clean, target in loader:
        timesteps = torch.randint(0, len(alpha), (clean.shape[0],), generator=generator)
        noise = torch.randn(clean.shape, generator=generator)
        a = alpha[timesteps].view(-1, 1, 1, 1)
        noisy = a.sqrt() * clean + (1.0 - a).sqrt() * noise
        clean_device = clean.to(device)
        target_device = target.to(device)
        noisy_logits = noisy_model(noisy.to(device), timesteps.to(device))
        clean_logits = clean_model(clean_device)
        noisy_correct += accuracy(noisy_logits, target_device)[0]
        clean_correct += accuracy(clean_logits, target_device)[0]
        noisy_loss += float(F.cross_entropy(noisy_logits, target_device, reduction="sum").cpu())
        clean_loss += float(F.cross_entropy(clean_logits, target_device, reduction="sum").cpu())
        total += target.numel()
    return {
        "test_examples": total,
        "noisy_random_t_accuracy": noisy_correct / total,
        "noisy_random_t_nll": noisy_loss / total,
        "clean_accuracy": clean_correct / total,
        "clean_nll": clean_loss / total,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=ROOT / "outputs" / "public-mnist-classifiers")
    parser.add_argument("--data-root", type=Path, default=ROOT / "outputs" / "public-mnist-classifiers" / "data")
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--learning-rate", type=float, default=3e-4)
    parser.add_argument("--weight-decay", type=float, default=1e-3)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = select_device(args.device)
    transform = transforms.Compose(
        [transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
    )
    train_set = datasets.MNIST(args.data_root, train=True, download=True, transform=transform)
    test_set = datasets.MNIST(args.data_root, train=False, download=True, transform=transform)
    loader_generator = torch.Generator().manual_seed(args.seed)
    train_loader = torch.utils.data.DataLoader(
        train_set,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=0,
        generator=loader_generator,
    )
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.batch_size, num_workers=0)

    scheduler = DDPMScheduler.from_pretrained(
        "Shiroi-max/ddpm-mnist-conditional-32",
        subfolder="scheduler",
        revision="368673bd7cab30e3e04af9ce5ecb1da2bf18f30f",
    )
    alpha = scheduler.alphas_cumprod.float()
    noisy_model = NoisyDigitClassifier().to(device)
    clean_model = CleanDigitClassifier().to(device)
    noisy_optimizer = torch.optim.AdamW(
        noisy_model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay
    )
    clean_optimizer = torch.optim.AdamW(
        clean_model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay
    )
    noisy_lr = torch.optim.lr_scheduler.CosineAnnealingLR(noisy_optimizer, T_max=args.epochs)
    clean_lr = torch.optim.lr_scheduler.CosineAnnealingLR(clean_optimizer, T_max=args.epochs)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    checkpoint_path = args.output_dir / "training_checkpoint.pt"
    start_epoch = 0
    best_noisy = best_clean = 0.0
    history: list[dict[str, float | int]] = []
    if args.resume and checkpoint_path.exists():
        state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        noisy_model.load_state_dict(state["noisy_model"])
        clean_model.load_state_dict(state["clean_model"])
        noisy_optimizer.load_state_dict(state["noisy_optimizer"])
        clean_optimizer.load_state_dict(state["clean_optimizer"])
        noisy_lr.load_state_dict(state["noisy_lr"])
        clean_lr.load_state_dict(state["clean_lr"])
        start_epoch = int(state["epoch"]) + 1
        best_noisy = float(state["best_noisy"])
        best_clean = float(state["best_clean"])
        history = list(state["history"])

    noise_generator = torch.Generator(device="cpu").manual_seed(args.seed + 1)
    for epoch in range(start_epoch, args.epochs):
        noisy_model.train()
        clean_model.train()
        noisy_seen = clean_seen = 0
        noisy_correct = clean_correct = 0
        noisy_loss_sum = clean_loss_sum = 0.0
        for clean, target in train_loader:
            timesteps = torch.randint(0, len(alpha), (clean.shape[0],), generator=noise_generator)
            noise = torch.randn(clean.shape, generator=noise_generator)
            a = alpha[timesteps].view(-1, 1, 1, 1)
            noisy = a.sqrt() * clean + (1.0 - a).sqrt() * noise
            clean_device = clean.to(device)
            target_device = target.to(device)
            noisy_optimizer.zero_grad(set_to_none=True)
            clean_optimizer.zero_grad(set_to_none=True)
            noisy_logits = noisy_model(noisy.to(device), timesteps.to(device))
            clean_logits = clean_model(clean_device)
            noisy_loss = F.cross_entropy(noisy_logits, target_device)
            clean_loss = F.cross_entropy(clean_logits, target_device)
            (noisy_loss + clean_loss).backward()
            noisy_optimizer.step()
            clean_optimizer.step()
            noisy_correct += accuracy(noisy_logits.detach(), target_device)[0]
            clean_correct += accuracy(clean_logits.detach(), target_device)[0]
            noisy_seen += target.numel()
            clean_seen += target.numel()
            noisy_loss_sum += float(noisy_loss.detach().cpu()) * target.numel()
            clean_loss_sum += float(clean_loss.detach().cpu()) * target.numel()
        noisy_lr.step()
        clean_lr.step()
        metrics = evaluate(noisy_model, clean_model, test_loader, alpha, device, args.seed + epoch)
        row: dict[str, float | int] = {
            "epoch": epoch,
            "train_examples": noisy_seen,
            "train_noisy_accuracy": noisy_correct / noisy_seen,
            "train_noisy_nll": noisy_loss_sum / noisy_seen,
            "train_clean_accuracy": clean_correct / clean_seen,
            "train_clean_nll": clean_loss_sum / clean_seen,
            **metrics,
        }
        history.append(row)
        best_noisy = max(best_noisy, float(metrics["noisy_random_t_accuracy"]))
        best_clean = max(best_clean, float(metrics["clean_accuracy"]))
        state = {
            "epoch": epoch,
            "noisy_model": noisy_model.state_dict(),
            "clean_model": clean_model.state_dict(),
            "noisy_optimizer": noisy_optimizer.state_dict(),
            "clean_optimizer": clean_optimizer.state_dict(),
            "noisy_lr": noisy_lr.state_dict(),
            "clean_lr": clean_lr.state_dict(),
            "best_noisy": best_noisy,
            "best_clean": best_clean,
            "history": history,
            "config": vars(args),
            "diffusion_repo": "Shiroi-max/ddpm-mnist-conditional-32",
            "diffusion_revision": "368673bd7cab30e3e04af9ce5ecb1da2bf18f30f",
        }
        torch.save(state, checkpoint_path)
        if float(metrics["noisy_random_t_accuracy"]) >= best_noisy:
            torch.save(noisy_model.state_dict(), args.output_dir / "noisy_classifier.pt")
        if float(metrics["clean_accuracy"]) >= best_clean:
            torch.save(clean_model.state_dict(), args.output_dir / "clean_judge.pt")
        with (args.output_dir / "metrics.csv").open("w", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=list(history[0]))
            writer.writeheader()
            writer.writerows(history)
        (args.output_dir / "run_metadata.json").write_text(
            json.dumps(
                {
                    "paper": "OAM1jJsMGp",
                    "scope": "auxiliary full-data real-model validation",
                    "device": str(device),
                    "train_examples_per_epoch": len(train_set),
                    "test_examples": len(test_set),
                    "epochs_requested": args.epochs,
                    "best_noisy_random_t_accuracy": best_noisy,
                    "best_clean_accuracy": best_clean,
                    "diffusion_repo": state["diffusion_repo"],
                    "diffusion_revision": state["diffusion_revision"],
                },
                indent=2,
            )
            + "\n"
        )
        print(json.dumps(row, sort_keys=True), flush=True)


if __name__ == "__main__":
    main()