"""CIFAR-10 Challenge! Your goal is to maximize test-set accuracy on CIFAR-10 (https://www.cs.toronto.edu/~kriz/cifar.html) using a feedforward (MLP) neural network before the end of class. Good luck! Once this script is running you can access the associated TensorBoard logs by running: tensorboard --logdir=runs Then open http://localhost:6006/ in a web browser. *** ADVICE FOR TUNING *** 1) Increase the capacity (width and/or depth) of your network until you start to see overfitting. 2) Tune your hyperparameters to address overfitting. Possible hyperparameters to explore: - Number of training epochs (track validation error) - Network size and shape - L2 regularization / Weight Decay - Dropout - Batch normalization 3) Change one thing at a time, and set `run_name` in main() to describe each experiment so you can tell your runs apart in TensorBoard. 4) Once you think you are doing as well as possible on the validation data, execute the testing code. IF YOU DO THIS MORE THAN ONCE, YOU ARE "CHEATING": tuning on the testing data invalidates your model evaluation. *** PYTORCH SYNTAX *** - Dropout: nn.Dropout(p=0.3) (a layer, usually placed after a ReLU) - Batch norm: nn.BatchNorm1d(num_units) (usually placed before a ReLU) - L2 penalty: torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4) *** TRAINING TIME *** Time per epoch is roughly proportional to the number of parameters printed when the script starts. A few million parameters trains in a few seconds per epoch on a laptop. Tens of millions will eat up your class time. *** EXPECTATIONS *** Well-tuned MLPs top out somewhere around 55-60% test accuracy on CIFAR-10. Don't be discouraged if your numbers look low! Note that, since this is image data, we could do significantly better by using image-specific approaches including convolutional neural networks and data augmentation. **Don’t do these things.** The point of this exercise is to practice tuning neural network hyperparameters, not to explore optimizations that are specific to a particular type of data. """ # References: # https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html # https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html # https://pytorch.org/tutorials/recipes/recipes/tensorboard_with_pytorch.html import time import numpy as np import torch import torchvision import torchvision.transforms as transforms from torch.utils.tensorboard import SummaryWriter import torch.nn as nn import matplotlib.pyplot as plt # Download CIFAR-10 from a JMU mirror; the official server is slow. torchvision.datasets.CIFAR10.url = ( "https://w3.cs.jmu.edu/spragunr/datasets/cifar-10-python.tar.gz" ) # ---------------------------- # Device selection # ---------------------------- device = (torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu") print(f"Using {device} device") # ---------------------------- # Reproducibility (optional but recommended) # ---------------------------- def seed_all(seed: int = 42): np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) seed_all(42) g = torch.Generator().manual_seed(42) # ---------------------------- # Helpers # ---------------------------- def count_parameters(model: nn.Module) -> int: return sum(p.numel() for p in model.parameters() if p.requires_grad) def plot_cifar_images(dataset, classes): """Show a 5x5 grid of (unnormalized) images from a normalized dataset. Assumes Normalize(mean=.5, std=.5) per channel; if you change normalization, adjust this unnormalize accordingly. """ plt.figure(figsize=(7, 7)) for i in range(25): img, label = dataset[i] # C x H x W tensor img = img / 2 + 0.5 # unnormalize for mean=.5, std=.5 npimg = img.numpy() plt.subplot(5, 5, i + 1) plt.title(classes[label], fontsize=8) plt.axis("off") plt.imshow(np.transpose(npimg, (1, 2, 0))) plt.suptitle("Sample CIFAR-10 training images with their labels\n" "(Close this window to start training.)") plt.tight_layout() plt.show() def plot_weights(weights: np.ndarray): """Plot first layer weights as images (assuming input is 3x32x32).""" num_show = min(100, weights.shape[0]) rows = int(np.ceil(np.sqrt(num_show))) plt.figure(figsize=(8, 8)) for i in range(num_show): img = weights[i, :].reshape(3, 32, 32).transpose((1, 2, 0)) mn, mx = np.min(img), np.max(img) rng = mx - mn img = (img - mn) / rng plt.subplot(rows, rows, i + 1) plt.axis("off") plt.imshow(img) plt.suptitle(f"First-layer weights for {num_show} of " f"{weights.shape[0]} hidden units, one image per unit\n" "(Close this window to end the script.)") plt.tight_layout() plt.show() def train(dataloader, model, loss_fn, optimizer): """Perform a single epoch of training.""" model.train() total_loss = 0.0 start_time = time.time() for batch, (X, y) in enumerate(dataloader): X, y = X.to(device), y.to(device) # Forward + loss pred = model(X) loss = loss_fn(pred, y) # Backprop optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() if batch % 100 == 0: print(".", end="", flush=True) print(f" Epoch time: {time.time() - start_time:.2f}s") return total_loss / len(dataloader) def evaluate(dataloader, model, loss_fn): """Evaluation helper.""" size = len(dataloader.dataset) num_batches = len(dataloader) model.eval() test_loss = 0.0 correct = 0.0 with torch.no_grad(): for X, y in dataloader: X = X.to(device) y = y.to(device) pred = model(X) test_loss += loss_fn(pred, y).item() correct += (pred.argmax(dim=1) == y).float().sum().item() return test_loss / num_batches, correct / size # ---------------------------- # Main experiment # ---------------------------- def main(): batch_size = 128 epochs = 10 learning_rate = 1e-3 run_name = "baseline" # Shows up in the TensorBoard run name transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) # ---------------------------- # Data # ---------------------------- train_set = torchvision.datasets.CIFAR10( root="./data", train=True, download=True, transform=transform ) train_set, val_set = torch.utils.data.random_split( train_set, [0.9, 0.1], generator=g ) pin = (device == "cuda") train_loader = torch.utils.data.DataLoader( train_set, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=pin, persistent_workers=True ) val_loader = torch.utils.data.DataLoader( val_set, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=pin, persistent_workers=True ) # Training accuracy is estimated from the first 5000 training # examples to save time. train_eval_loader = torch.utils.data.DataLoader( torch.utils.data.Subset(train_set, range(5000)), batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=pin, persistent_workers=True ) classes = ( 'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck' ) # Show some sample images plot_cifar_images(train_set, classes) # ---------------------------- # Model # ---------------------------- # Starter MLP (students should widen/deepen, add dropout/bn, etc.) model = nn.Sequential( nn.Flatten(), nn.Linear(32 * 32 * 3, 16), nn.ReLU(), nn.Linear(16, 10), ).to(device) print(model) print("Total parameters:", count_parameters(model)) loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.Adam( model.parameters(), lr=learning_rate ) # ---------------------------- # TensorBoard # ---------------------------- writer = SummaryWriter(comment="_" + run_name) imgs, _ = next(iter(train_loader)) writer.add_graph(model, imgs.to(device)) # ---------------------------- # Training loop # ---------------------------- for epoch in range(1, epochs + 1): print(f"Epoch {epoch}/{epochs}") train_loss_epoch = train(train_loader, model, loss_fn, optimizer) train_loss, train_acc = evaluate(train_eval_loader, model, loss_fn) val_loss, val_acc = evaluate(val_loader, model, loss_fn) print( f" Train loss: {train_loss:.4f} acc: {train_acc:.3f} | " f"Val loss: {val_loss:.4f} acc: {val_acc:.3f}" ) # Store data for TensorBoard visualization writer.add_scalar("Loss/train", train_loss, epoch) writer.add_scalar("Loss/val", val_loss, epoch) writer.add_scalar("Accuracy/train", train_acc, epoch) writer.add_scalar("Accuracy/val", val_acc, epoch) torch.save(model, "latest_model.pt") # ---------------------------- # Visualize first-layer weights # ---------------------------- wts = model[1].weight.detach().cpu().numpy() plot_weights(wts) writer.close() if __name__ == "__main__": main()