Threat Models, Certified Defences and Backdoor Attacks

def add_square_trigger_numpy(img: np.ndarray, size: int = 3, value: int = 255, location: str = "bottom_left", pad: int = 1):
    """
    Add a black/white square trigger to a (1,H,W) numpy array in [0,255] uint8.
    - size: side length in pixels
    - value: pixel intensity (0=black, 255=white)
    - location: position of the square ('bottom_left' etc.)
    - pad: offset from the edge
    Returns a modified copy of the image.
    """
    assert img.ndim == 3 and img.shape[0] == 1, "expected (1,H,W)"
    _, H, W = img.shape
    arr = img.copy()

    if location == "bottom_left":
        r0, c0 = H - pad - size, pad
    elif location == "bottom_right":
        r0, c0 = H - pad - size, W - pad - size
    elif location == "top_left":
        r0, c0 = pad, pad
    elif location == "top_right":
        r0, c0 = pad, W - pad - size
    else:
        raise ValueError(f"Unknown location: {location}")

    r1, c1 = r0 + size, c0 + size
    arr[:, r0:r1, c0:c1] = value
    return arr

poisoned_imgs = []
for i in range(len(imgs)):
    poisoned_imgs.append(torch.from_numpy(add_square_trigger_numpy(imgs[i].numpy(), size=3, value=255, location="bottom_left", pad=1)))
show_image_grid(list(poisoned_imgs), nrow=10)

backdoored_test_imgs = []
for i in range(10):
    backdoored_test_imgs.append(mnist_test_poisoned.data[i].unsqueeze(0))
show_image_grid(list(backdoored_test_imgs), nrow=10)
model.eval()
pred = model(torch.stack(backdoored_test_imgs).to(device)/255.0).argmax(dim=1)
print("Predicted labels for backdoored test samples:", pred)
# you should see all predicted as target label 0

# Cell: InMemoryPoisonedMNIST - subclass of torchvision.datasets.MNIST
from typing import Optional, Dict, Any, Sequence
import os
import json
import random
import numpy as np
import torch
from torchvision.datasets import MNIST

class InMemoryPoisonedMNIST(MNIST):
    """
    Subclass of torchvision.datasets.MNIST that can apply an *in-memory*
    deterministic data-poisoning modification to the dataset tensors.

    After calling `apply_poison(...)` the underlying `self.data` and `self.targets`
    are modified (uint8 images in `self.data` are overwritten in-place).
    This simulates a dataset that was downloaded already poisoned.

    Parameters
    ----------
    root : str
        dataset root passed to MNIST
    train : bool
        train or test split (poisoning typically applied to train=True)
    download : bool
        whether to download
    """
    @property
    def raw_folder(self) -> str:
        # Reuse ./data/MNIST/raw (by default torchvision would download a second copy
        # into ./data/InMemoryPoisonedMNIST/raw, named after this class)
        return os.path.join(self.root, "MNIST", "raw")

    def __init__(self, root: str = "./data", train: bool = True, download: bool = True, **kwargs):
        super().__init__(root=root, train=train, download=download, **kwargs)
        # internal bookkeeping
        self.poisoned_indices: Optional[Sequence[int]] = None
        self.poison_metadata: Optional[Dict[str, Any]] = None
        # keep a small in-memory copy if user wants to be able to restore later
        self._orig_data = None
        self._orig_targets = None

    def _compute_coords(self, H: int, W: int, size: int, location: str, pad: int):
        """Helper: compute (r0,r1,c0,c1) for square given image HxW."""
        s = int(size)
        if location == "bottom_left":
            r0, c0 = H - pad - s, pad
        elif location == "bottom_right":
            r0, c0 = H - pad - s, W - pad - s
        elif location == "top_left":
            r0, c0 = pad, pad
        elif location == "top_right":
            r0, c0 = pad, W - pad - s
        elif location == "center":
            r0, c0 = (H - s) // 2, (W - s) // 2
        else:
            raise ValueError(f"Unknown location: {location}")
        r1, c1 = r0 + s, c0 + s
        # clamp
        r0, r1 = max(0, r0), min(H, r1)
        c0, c1 = max(0, c0), min(W, c1)
        return r0, r1, c0, c1

    def apply_poison(
        self,
        p_poison: float = 0.02,
        target_label: int = 0,
        size: int = 3,
        value: int = 255,
        location: str = "bottom_left",
        pad: int = 1,
        seed: int = 42,
        overwrite: bool = True,  # kept for API compatibility (no-op here)
        save_indices_path: Optional[str] = None,
        save_metadata_path: Optional[str] = None,
        keep_original_copy: bool = False,
    ) -> Dict[str, Any]:
        """
        Apply in-memory poisoning using add_square_trigger_numpy on (1,H,W) arrays.
        """
        assert 0.0 <= p_poison <= 1.0, "p_poison must be in [0,1]"
        n = len(self.data)
        n_poison = int(round(p_poison * n))
        rng = random.Random(seed)
        poisoned_idx = sorted(rng.sample(range(n), n_poison))

        # optional backup
        if keep_original_copy and (self._orig_data is None):
            self._orig_data = self.data.clone()
            self._orig_targets = self.targets.clone()

        # ensure tensor types
        if not isinstance(self.data, torch.Tensor):
            self.data = torch.tensor(self.data, dtype=torch.uint8)
        else:
            # work on CPU uint8
            if self.data.dtype != torch.uint8:
                self.data = (self.data.clone() * 255).to(torch.uint8)
            if self.data.device.type != "cpu":
                self.data = self.data.cpu()

        if not isinstance(self.targets, torch.Tensor):
            self.targets = torch.tensor(self.targets)

        # apply trigger via numpy helper per selected index
        for idx in poisoned_idx:
            # self.data[idx]: [H, W] uint8
            hw = self.data[idx].numpy()                                # (H, W), uint8 (shared mem view)
            chw = hw[None, ...]                                        # (1, H, W)
            poisoned = add_square_trigger_numpy(
                chw, size=size, value=value, location=location, pad=pad
            )                                                          # (1, H, W), uint8
            self.data[idx] = torch.from_numpy(poisoned[0])             # back to [H, W] uint8
            self.targets[idx] = int(target_label)

        metadata = {
            "n_total": n,
            "n_poison": n_poison,
            "p_poison": float(p_poison),
            "target_label": int(target_label),
            "trigger": {"size": int(size), "value": int(value), "location": location, "pad": int(pad)},
            "seed": int(seed),
            "poisoned_indices": poisoned_idx,
        }
        self.poisoned_indices = poisoned_idx
        self.poison_metadata = metadata

        if save_indices_path is not None:
            np.save(save_indices_path, np.array(poisoned_idx, dtype=np.int32))
        if save_metadata_path is not None:
            with open(save_metadata_path, "w") as f:
                json.dump(metadata, f, indent=2)

        return metadata

    def restore_original(self):
        """
        Restore original data/targets if keep_original_copy=True was used.
        """
        if self._orig_data is None or self._orig_targets is None:
            raise RuntimeError("No original copy stored. Create with keep_original_copy=True in apply_poison().")
        self.data = self._orig_data.clone()
        self.targets = self._orig_targets.clone()
        self.poisoned_indices = None
        self.poison_metadata = None

# -------------------------
# Example usage
# -------------------------
root = "./data"
poisoned_mnist = InMemoryPoisonedMNIST(root=root, train=True, download=True, transform=transform)

# Apply a 2% poisoning: bottom-left 3x3 white square -> target label 0
meta = poisoned_mnist.apply_poison(
    p_poison=0.02,
    target_label=0,
    size=3,
    value=255,
    location="bottom_left",
    pad=1,
    seed=42,
    save_indices_path="./data/poisoned_indices.npy",
    save_metadata_path="./data/poison_metadata.json",
    keep_original_copy=True,   # optional: allow restore if you want
)

poisoned_imgs = []
for idx in poisoned_mnist.poisoned_indices[:10]:
    poisoned_imgs.append(poisoned_mnist.data[idx].unsqueeze(0))
show_image_grid(list(poisoned_imgs), nrow=10)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from typing import Tuple, Optional

# -------------------------
# LeNet-5 style model
# -------------------------
class LeNet(nn.Module):
    """
    LeNet-5 style convnet adapted for MNIST (1x28x28 input).
    Architecture (classic):
      Conv1(1->6, kernel=5) -> ReLU -> AvgPool(2)
      Conv2(6->16, kernel=5) -> ReLU -> AvgPool(2)
      Flatten -> FC(16*4*4 -> 120) -> ReLU
             -> FC(120 -> 84) -> ReLU
             -> FC(84 -> num_classes)
    """
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)      # 28 -> 24
        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)  # 24 -> 12
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)     # 12 -> 8
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)  # 8 -> 4

        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

        # optional small init (closer to original LeNet)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(x.shape[0], -1)
        x = F.relu(self.fc1(x))
        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        x = F.relu(self.fc2(x))
        x = self.fc3(x)  # logits
        return x

from tqdm import tqdm

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
        self.max = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
        self.max = max(self.max, val)

def accuracy(output, target, topk=(1,)):
    maxk = max(topk)

    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].float().sum()
        res.append(correct_k.mul_(1/batch_size))
    return res

@torch.no_grad()
def evaluate(model, eval_data_loader):
    model.eval()
    # Training Evaluations
    loss_meters = AverageMeter()
    acc_meters = AverageMeter()
    loss_list, correct_list = [], []
    for i, data in enumerate(eval_data_loader):
        # Prepare batch data
        images, labels = data
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
        batch_size = images.shape[0]
        logits = model(images)
        loss = F.cross_entropy(logits, labels, reduction='none')
        loss_list += loss.detach().cpu().numpy().tolist()
        loss = loss.mean().item()
        # Calculate acc
        acc = accuracy(logits, labels, topk=(1,))[0].item()

        # Update Meters
        loss_meters.update(loss, batch_size)
        acc_meters.update(acc, batch_size)

        _, predicted = torch.max(logits.data, 1)
        correct = (predicted == labels)
        correct_list += correct.detach().cpu().numpy().tolist()
    stats = {   
        "eval_loss": loss_meters.avg,
        "eval_acc": acc_meters.avg,
        "eval_loss_list": loss_list,
        "eval_correct_list": correct_list,
    }
    return stats


def train(model, optimizer, train_loader):

    # Set Meters
    loss_meters = AverageMeter()
    acc_meters = AverageMeter()

    # Training
    model.train()
    pbar = tqdm(train_loader)
    for i, data in enumerate(pbar):
        # Prepare batch data
        images, labels = data
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
        batch_size = images.shape[0]
        model.zero_grad()
        optimizer.zero_grad()

        # Objective function
        logits = model(images)
        loss = F.cross_entropy(logits, labels)
        
        # Optimize
        loss.backward()
        optimizer.step()
        # Calculate acc
        loss = loss.item()
        acc = accuracy(logits, labels, topk=(1,))[0].item()

        # Update Meters
        loss_meters.update(loss, batch_size)
        acc_meters.update(acc, batch_size)
        pbar.set_postfix(loss=loss, acc=acc*100)

    stats = {
        "train_loss": loss_meters.avg,
        "train_acc": acc_meters.avg,
    }

    return stats

# ---------- Hyperparams ----------
root = "./data"
epochs = 5
lr = 0.01
momentum = 0.9
weight_decay = 5e-4
num_workers = 0
# ---------------------------------
model = LeNet(num_classes=10).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)

# Triggered test set for the attack success rate (ASR):
# every test image gets the bottom-left 3x3 white square and the label 0
mnist_test_poisoned = InMemoryPoisonedMNIST(root=root, train=False, download=True, transform=transform)
meta_test = mnist_test_poisoned.apply_poison(
    p_poison=1.0,
    target_label=0,
    size=3,
    value=255,
    location="bottom_left",
    pad=1,
)
# Drop the digits that are already 0: predicting 0 for them is not a success of the attack
non_target = mnist_test.targets != 0
mnist_test_poisoned.data = mnist_test_poisoned.data[non_target]
mnist_test_poisoned.targets = mnist_test_poisoned.targets[non_target]

train_loader = DataLoader(poisoned_mnist, batch_size=64, shuffle=True, num_workers=num_workers)
test_loader = DataLoader(mnist_test, batch_size=256, shuffle=False, num_workers=num_workers)
test_poisoned_loader = DataLoader(mnist_test_poisoned, batch_size=256, shuffle=False, num_workers=num_workers)

# ------------------------- training loop -------------------------
for epoch in range(1, epochs + 1):
    print(f"Epoch {epoch}/{epochs}")
    train_stats = train(model, optimizer, train_loader)
    eval_stats_clean = evaluate(model, test_loader)
    eval_stats_poisoned = evaluate(model, test_poisoned_loader)

    print(f"  Train                  - loss: {train_stats['train_loss']:.4f}, acc: {train_stats['train_acc']*100:.2f}%")
    print(f"  Clean test set         - CA:  {eval_stats_clean['eval_acc']*100:.2f}%")
    print(f"  Triggered test (non-0) - ASR: {eval_stats_poisoned['eval_acc']*100:.2f}%")
    print("-" * 40)
# Expect both clean accuracy and attack success rate around 98-99% after 5 epochs

from sklearn.manifold import TSNE

# get all indices and subtract poisoned ones
all_idx = np.arange(len(poisoned_mnist))
poisoned_idx = set(poisoned_mnist.poisoned_indices)
clean_candidates = [i for i in all_idx if i not in poisoned_idx]

poisoned_imgs = [poisoned_mnist.data[idx] for idx in poisoned_idx]
poisoned_imgs = torch.stack(poisoned_imgs).unsqueeze(1)  # (N, 1, H, W)
poisoned_imgs = poisoned_imgs.to(device) / 255.0  # normalize to [0,1]

clean_imgs = [poisoned_mnist.data[idx] for idx in clean_candidates[:len(poisoned_idx)]]
clean_imgs = torch.stack(clean_imgs).unsqueeze(1)  # (N, 1, H, W)
clean_imgs = clean_imgs.to(device) / 255.0  # normalize to [0,1]

model.eval()
with torch.no_grad():
    poisoned_feats = model.forward_features(poisoned_imgs).cpu()
    clean_feats = model.forward_features(clean_imgs).cpu()

# Concatenate and run t-SNE
X = torch.cat([clean_feats, poisoned_feats], dim=0).numpy()        # [N, D]
is_poison = np.array([0]*len(clean_feats) + [1]*len(poisoned_feats), dtype=int)

tsne = TSNE(n_components=2, perplexity=30, max_iter=1000,
            learning_rate=200, init="pca", random_state=42, verbose=1)
X2 = tsne.fit_transform(X)

# Plot
plt.figure(figsize=(7.5, 6))
plt.scatter(X2[is_poison==0, 0], X2[is_poison==0, 1],
            s=10, marker='o', facecolors='none', edgecolors='#1f77b4', linewidths=0.8, label='Clean')
plt.scatter(X2[is_poison==1, 0], X2[is_poison==1, 1],
            s=12, marker='x', c='#d62728', label='Triggered')

plt.title("t-SNE of LeNet penultimate features\nClean vs Triggered ")
plt.xticks([]); plt.yticks([]); plt.legend(frameon=False)
plt.tight_layout()
plt.show()

# -------------------------
# k-NN distance outlier score
# -------------------------
@torch.no_grad()
def knn_score(f: torch.Tensor, k: int = 16, chunk_size: int = 2000) -> torch.Tensor:
    """
    Mean L2 distance from each row of `f` (N x D) to its k nearest
    neighbours among all N rows, excluding itself.

    Args:
        f: [N, D] feature tensor (float). Put it on the desired device.
        k: number of nearest neighbours to average.
        chunk_size: rows per distance block, so the N x N matrix is never stored.

    Returns:
        score: [N] tensor, higher means more outlier-like.

    The lecture's SLOF goes one step further and divides k-dist(q)
    by the k-dist of q's neighbours.
    """
    assert f.dim() == 2, "f must be [N, D]"
    N = f.size(0)
    k = max(1, min(k, N - 1))
    score = torch.empty(N, device=f.device)
    for i in range(0, N, chunk_size):
        d = torch.cdist(f[i:i + chunk_size], f, p=2)               # [chunk, N]
        rows = torch.arange(d.size(0), device=f.device)
        d[rows, rows + i] = float("inf")                           # remove self-distance
        nearest = torch.topk(d, k, dim=1, largest=False).values    # [chunk, k]
        score[i:i + chunk_size] = nearest.mean(dim=1)
    return score

eval_train_loader = DataLoader(poisoned_mnist, batch_size=256, shuffle=False, num_workers=num_workers)
model.eval()
with torch.no_grad():
    feats = torch.cat([model.forward_features(images.to(device)) for images, _ in eval_train_loader])

# neighbours are searched over all 60 000 training samples
scores = knn_score(feats, k=16).cpu().numpy()
print(f"Computed k-NN scores for {len(scores)} training samples.")

from sklearn.metrics import roc_auc_score

y = np.zeros(len(poisoned_mnist), dtype=int)
y[poisoned_mnist.poisoned_indices] = 1  # 1 = poisoned
# ROC AUC
roc_auc = roc_auc_score(y, scores)
print(f"ROC AUC: {roc_auc*100:.2f}%")

def total_variation_loss(img, weight=1):
    b, c, h, w = img.size()
    tv_h = torch.pow(img[:, :, 1:, :]-img[:, :, :-1, :], 2).sum(dim=[1, 2, 3])
    tv_w = torch.pow(img[:, :, :, 1:]-img[:, :, :, :-1], 2).sum(dim=[1, 2, 3])
    return weight*(tv_h+tv_w)/(c*h*w)


class CognitiveDistillation(nn.Module):
    def __init__(self, lr=0.1, p=1, gamma=0.01, beta=1.0, num_steps=100, mask_channel=1, norm_only=False):
        super(CognitiveDistillation, self).__init__()
        self.p = p
        self.gamma = gamma
        self.beta = beta
        self.num_steps = num_steps
        self.l1 = torch.nn.L1Loss(reduction='none')
        self.lr = lr
        self.mask_channel = mask_channel
        self.get_features = False
        self._EPSILON = 1.e-6
        self.norm_only = norm_only

    def get_raw_mask(self, mask):
        mask = (torch.tanh(mask) + 1) / 2
        return mask

    def forward(self, model, images):
        model.eval()
        if torch.min(images) < 0 or torch.max(images) > 1:
            raise ValueError("images should be in [0, 1]")
        b, c, h, w = images.shape
        mask = torch.ones(b, self.mask_channel, h, w).to(images.device)
        mask_param = nn.Parameter(mask)
        optimizerR = torch.optim.Adam([mask_param], lr=self.lr, betas=(0.1, 0.1))
        logits = model(images).detach()
        for step in range(self.num_steps):
            optimizerR.zero_grad()
            mask = self.get_raw_mask(mask_param).to(images.device)
            x_cp = images * mask + (1-mask) * torch.rand(b, c, 1, 1).to(images.device)
            adv_logits = model(x_cp)
            loss = self.l1(adv_logits, logits).mean(dim=1)
            norm = torch.norm(mask, p=self.p, dim=[1, 2, 3])
            norm = norm * self.gamma
            loss_total = loss + norm + self.beta * total_variation_loss(mask)
            loss_total.mean().backward()
            optimizerR.step()
        mask = self.get_raw_mask(mask_param).detach().cpu()
        return mask.detach()

cognitive_distiller = CognitiveDistillation(lr=0.1, p=1, gamma=0.05, beta=1.0, num_steps=200, mask_channel=1)
# Get some poisoned samples
poisoned_samples = []
for idx in poisoned_mnist.poisoned_indices[:10]:
    poisoned_samples.append(poisoned_mnist.data[idx].unsqueeze(0))
poisoned_samples = torch.stack(poisoned_samples).to(device) / 255.0  # (N, 1, H, W) in [0,1]
# Run cognitive distillation to get masks
masks = cognitive_distiller(model, poisoned_samples)  # (N, 1, H, W) in [0,1]
# Show poisoned samples
show_image_grid(list(poisoned_samples), nrow=10)
# Show masks
show_image_grid(list(masks), nrow=10)
# Overlay masks on images
overlayed = poisoned_samples.cpu() * masks + (1 - masks) * 0.
show_image_grid(list(overlayed), nrow=10)

# All 1200 poisoned images and 3000 random clean ones (all 60 000 is slow on a CPU)
rng = np.random.RandomState(0)
clean_pool = np.setdiff1d(np.arange(len(poisoned_mnist)), poisoned_mnist.poisoned_indices)
clean_idx = rng.choice(clean_pool, 3000, replace=False)
subset_idx = np.concatenate([poisoned_mnist.poisoned_indices, clean_idx])
subset_imgs = poisoned_mnist.data[subset_idx].unsqueeze(1).to(device) / 255.0  # [0,1]

model.eval()
masks = torch.cat([cognitive_distiller(model, subset_imgs[i:i + 600])
                   for i in range(0, len(subset_imgs), 600)])
cd_scores = masks.norm(p=1, dim=[1, 2, 3]).numpy()  # L1 norm of each mask

print(f"Computed Cognitive Distillation scores for {len(cd_scores)} training samples.")

y_sub = np.r_[np.ones(len(poisoned_mnist.poisoned_indices)), np.zeros(len(clean_idx))].astype(int)
# With gamma = 0.05 the masks of clean images shrink to almost nothing,
# while those of poisoned images keep the trigger: a large mask flags a poison.
print(f"Mean mask L1 norm: poisoned {cd_scores[y_sub == 1].mean():.1f}, clean {cd_scores[y_sub == 0].mean():.1f}")
roc_auc = roc_auc_score(y_sub, cd_scores)
print(f"ROC AUC: {roc_auc*100:.2f}%")

! pip install torch torchvision matplotlib numpy

# Notebook cell: imports, reproducible seed, device, and MNIST download + poisoned/triggered dataset wrappers

import os
import random
from typing import Optional, Tuple

import numpy as np
from PIL import Image
import torch
from torch import Tensor
from torch.utils.data import Dataset
import torchvision
from torchvision import transforms
from torchvision.datasets import MNIST

# -------------------------
# Reproducibility & device
# -------------------------
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

def get_device() -> torch.device:
    """
    Return the best available device: CUDA > MPS > CPU.
    """
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() and torch.backends.mps.is_built():
        return torch.device("mps")
    return torch.device("cpu")

device = get_device()
print("Using device:", device)

# -------------------------
# Root directory for data
# -------------------------
root = "./data"
os.makedirs(root, exist_ok=True)

# -------------------------
# Define basic transform
# (convert to tensor)
# -------------------------
transform = transforms.Compose([
    transforms.PILToTensor(),                         # (1, H, W) uint8
    transforms.ConvertImageDtype(torch.float32),      # -> float32 in [0,1]
])

# -------------------------
# Download and create datasets
# -------------------------
mnist_train = MNIST(root=root, train=True, download=True, transform=transform)
mnist_test = MNIST(root=root, train=False, download=True, transform=transform)

print(f"Training set size: {len(mnist_train)}")
print(f"Test set size: {len(mnist_test)}")

# Example: show one sample (tensor shape + label)
img, label = mnist_train[0]
print("Sample image shape:", img.shape)
print("Sample label:", label)

import matplotlib.pyplot as plt
%matplotlib inline

def show_image_grid(imgs, nrow=10):
    grid_img = torchvision.utils.make_grid(imgs, nrow=nrow)
    fig= plt.figure(figsize=(10, 500), dpi=80)
    plt.imshow(grid_img.permute(1, 2, 0).detach().cpu().numpy())
    plt.show()
    plt.close()

# ------------------------
# Randomly pick some indices
# ------------------------
n_samples = 20
idxs = np.random.choice(len(mnist_train), size=n_samples, replace=False)

imgs, labels = zip(*[mnist_train[i] for i in idxs])

print("Labels:", labels)
show_image_grid(list(imgs), nrow=10)