"""Benchmark the MNIST MLP with PyTorch on CPU or CUDA.

Install: pip install torch torchvision numpy
Run: python mnist-pytorch.py --device cpu
Profile: nsys profile -o pytorch-profile python mnist-pytorch.py --device cuda

Times one training epoch plus validation, including data loading from in-memory
arrays, CPU-to-device transfers and metrics. Includes partial batches. Download,
initialization and warmup are excluded. These are new measurements; they do not
reproduce the historical framework timings in the article.
"""

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision.datasets import MNIST
import numpy as np

def run_epoch(model, X, y, batch_size=32, optimizer=None):
    training = optimizer is not None
    model.train(training)
    device = next(model.parameters()).device
    loader = DataLoader(
        TensorDataset(torch.as_tensor(X, dtype=torch.float32),
                      torch.as_tensor(y, dtype=torch.long)),
        batch_size=batch_size, shuffle=training,
    )
    loss_fn = nn.CrossEntropyLoss()
    total_loss, correct = 0.0, 0
    with torch.enable_grad() if training else torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            if training:
                optimizer.zero_grad()
            logits = model(xb)
            loss = loss_fn(logits, yb)
            if training:
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * len(yb)
            correct += (logits.argmax(dim=1) == yb).sum().item()
    return total_loss / len(y), correct / len(y)


def benchmark(device="cpu", batch_size=32, epochs=1, samples=60000, threads=1):
    from copy import deepcopy
    from time import perf_counter

    torch.set_num_threads(threads)
    torch.manual_seed(42)
    device = torch.device(device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA is unavailable; use --device cpu")
    train_data = MNIST("./data", train=True, download=True)
    X = train_data.data[:samples].reshape(-1, 784).float() / 255.0
    y = train_data.targets[:samples]
    split = int(0.8 * len(X))
    model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10)).to(device)
    initial = deepcopy(model.state_dict())
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    # Warm up outside the timing, then restore the starting weights and RNG.
    run_epoch(model, X[:batch_size], y[:batch_size], batch_size, optimizer)
    model.load_state_dict(initial)
    torch.manual_seed(42)
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    start = perf_counter()
    for _ in range(epochs):
        loss, accuracy = run_epoch(model, X[:split], y[:split], batch_size, optimizer)
        val_loss, val_accuracy = run_epoch(model, X[split:], y[split:], batch_size)
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elapsed = perf_counter() - start
    result = dict(device=str(device), batch_size=batch_size, epochs=epochs,
                  training_images=split, validation_images=len(X)-split,
                  threads=threads, seconds_per_epoch=elapsed / epochs,
                  loss=loss, accuracy=accuracy, val_loss=val_loss, val_accuracy=val_accuracy,
                  torch_version=torch.__version__)
    print(result)
    return result


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--samples", type=int, default=60000)
    parser.add_argument("--threads", type=int, default=1)
    args = parser.parse_args()
    if args.batch_size < 1 or args.epochs < 1 or not 2 <= args.samples <= 60000 or args.threads < 1:
        parser.error("Require positive batch size, epochs and threads, and 2–60000 samples")
    benchmark(args.device, args.batch_size, args.epochs, args.samples, args.threads)