Conv Nets Py Torch

import torch
import torch.nn as nn
import torchvision
import matplotlib.pyplot as plt
from torchinfo import summary
import numpy as np

# Self-contained course PyTorch computer-vision helpers.
from torchvision.transforms import ToTensor
from sklearn.datasets import load_digits
from PIL import Image
import glob

default_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_mnist(batch_size=64):
    global data_train, data_test, train_loader, test_loader
    course_digits = load_digits()
    images = torch.as_tensor(course_digits.images, dtype=torch.float32).unsqueeze(1) / 16.0
    images = torch.nn.functional.interpolate(images, size=(28, 28), mode="bilinear", align_corners=False)
    labels = torch.as_tensor(course_digits.target, dtype=torch.long)
    split = int(0.8 * len(images))
    data_train = torch.utils.data.TensorDataset(images[:split], labels[:split])
    data_test = torch.utils.data.TensorDataset(images[split:], labels[split:])
    train_loader = torch.utils.data.DataLoader(data_train, batch_size=batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(data_test, batch_size=batch_size)
    return data_train, data_test

def _validate(net, dataloader, loss_fn):
    net.eval()
    total_loss, correct, count = 0.0, 0, 0
    with torch.no_grad():
        for features, labels in dataloader:
            features, labels = features.to(default_device), labels.to(default_device)
            output = net(features)
            total_loss += loss_fn(output, labels).item() * len(labels)
            correct += (output.argmax(1) == labels).sum().item()
            count += len(labels)
    return total_loss / max(count, 1), correct / max(count, 1)

def train(net, train_loader, test_loader, optimizer=None, lr=0.01, epochs=10, loss_fn=None):
    net.to(default_device)
    loss_fn = loss_fn or nn.CrossEntropyLoss()
    optimizer = optimizer or torch.optim.Adam(net.parameters(), lr=lr)
    history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
    for epoch in range(epochs):
        net.train()
        total_loss, correct, count = 0.0, 0, 0
        for features, labels in train_loader:
            features, labels = features.to(default_device), labels.to(default_device)
            optimizer.zero_grad()
            output = net(features)
            loss = loss_fn(output, labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * len(labels)
            correct += (output.argmax(1) == labels).sum().item()
            count += len(labels)
        validation_loss, validation_accuracy = _validate(net, test_loader, loss_fn)
        history["train_loss"].append(total_loss / max(count, 1))
        history["train_acc"].append(correct / max(count, 1))
        history["val_loss"].append(validation_loss)
        history["val_acc"].append(validation_accuracy)
        print(f"Epoch {epoch + 1}: train_acc={history['train_acc'][-1]:.3f}, val_acc={validation_accuracy:.3f}")
    return history

def train_long(net, train_loader, test_loader, **kwargs):
    return train(net, train_loader, test_loader, **kwargs)

def plot_results(history):
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    axes[0].plot(history["train_acc"], label="train")
    axes[0].plot(history["val_acc"], label="validation")
    axes[1].plot(history["train_loss"], label="train")
    axes[1].plot(history["val_loss"], label="validation")
    axes[0].set_title("Accuracy")
    axes[1].set_title("Loss")
    axes[0].legend()
    axes[1].legend()
    plt.show()

def plot_convolution(kernel, title=""):
    convolution = nn.Conv2d(1, 1, kernel_size=3, bias=False)
    with torch.no_grad():
        convolution.weight.copy_(kernel.reshape(1, 1, 3, 3))
        fig, axes = plt.subplots(2, 6, figsize=(10, 4))
        fig.suptitle(title)
        for index in range(5):
            image = data_train[index][0]
            axes[0, index].imshow(image[0], cmap="gray")
            axes[1, index].imshow(convolution(image.unsqueeze(0))[0, 0], cmap="gray")
            axes[0, index].axis("off")
            axes[1, index].axis("off")
        axes[0, 5].imshow(kernel, cmap="gray")
        axes[0, 5].axis("off")
        axes[1, 5].axis("off")
        plt.show()

def display_dataset(dataset, n=10, classes=None):
    count = min(n, len(dataset))
    fig, axes = plt.subplots(1, count, figsize=(1.8 * count, 3))
    axes = np.atleast_1d(axes)
    for index in range(count):
        image, label = dataset[index]
        image = image.detach().cpu()
        low, high = image.min(), image.max()
        image = (image - low) / (high - low + 1e-7)
        axes[index].imshow(np.transpose(image.numpy(), (1, 2, 0)))
        axes[index].axis("off")
        if classes is not None:
            axes[index].set_title(classes[int(label)])
    plt.show()

def check_image_dir(pattern):
    invalid = []
    for filename in glob.glob(pattern):
        try:
            with Image.open(filename) as image:
                image.verify()
        except (OSError, SyntaxError):
            invalid.append(filename)
    if invalid:
        raise ValueError(f"Invalid images detected: {invalid[:5]}")
    return 0
load_mnist(batch_size=128)

plot_convolution(torch.tensor([[-1.,0.,1.],[-1.,0.,1.],[-1.,0.,1.]]),'Vertical edge filter')
plot_convolution(torch.tensor([[-1.,-1.,-1.],[0.,0.,0.],[1.,1.,1.]]),'Horizontal edge filter')

class OneConv(nn.Module):
    def __init__(self):
        super(OneConv, self).__init__()
        self.conv = nn.Conv2d(in_channels=1,out_channels=9,kernel_size=(5,5))
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(5184,10)

    def forward(self, x):
        x = nn.functional.relu(self.conv(x))
        x = self.flatten(x)
        x = nn.functional.log_softmax(self.fc(x),dim=1)
        return x

net = OneConv()

summary(net,input_size=(1,1,28,28))

hist = train(net,train_loader,test_loader,epochs=5)
plot_results(hist)

fig,ax = plt.subplots(1,9)
with torch.no_grad():
    p = next(net.conv.parameters())
    for i,x in enumerate(p):
        ax[i].imshow(x.detach().cpu()[0,...])
        ax[i].axis('off')

class MultiLayerCNN(nn.Module):
    def __init__(self):
        super(MultiLayerCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(10, 20, 5)
        self.fc = nn.Linear(320,10)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 320)
        x = nn.functional.log_softmax(self.fc(x),dim=1)
        return x

net = MultiLayerCNN()
summary(net,input_size=(1,1,28,28))

hist = train(net,train_loader,test_loader,epochs=5)

# course-edition deterministic RGB pattern fixture v1
_course_generator = torch.Generator().manual_seed(2026)
_course_axis = torch.linspace(-1.0, 1.0, 32)
_course_y, _course_x = torch.meshgrid(_course_axis, _course_axis, indexing="ij")

def _course_rgb_pattern(label):
    frequency = 1 + label % 3
    phase = label * 0.37
    horizontal = 0.5 + 0.5 * torch.sin(frequency * torch.pi * _course_x + phase)
    vertical = 0.5 + 0.5 * torch.cos((1 + label // 3) * torch.pi * _course_y - phase)
    radius = torch.sqrt((_course_x - 0.12 * (label % 3 - 1)) ** 2 + (_course_y - 0.12 * (label // 3 - 1)) ** 2)
    ring = torch.exp(-((radius - (0.25 + 0.035 * label)) ** 2) / 0.012)
    channels = torch.stack((horizontal, vertical, ring))
    return torch.roll(channels, shifts=label % 3, dims=0).clamp(0.0, 1.0)

_course_templates = torch.stack([_course_rgb_pattern(label) for label in range(10)])
_course_labels = torch.arange(500, dtype=torch.long) % 10
_course_images = _course_templates[_course_labels].clone()
_course_noise = 0.04 * torch.rand(_course_images.shape, generator=_course_generator)
_course_images = ((_course_images + _course_noise).clamp(0.0, 1.0) - 0.5) / 0.5
_course_split = 400
trainset = torch.utils.data.TensorDataset(_course_images[:_course_split], _course_labels[:_course_split])
testset = torch.utils.data.TensorDataset(_course_images[_course_split:], _course_labels[_course_split:])
trainloader = torch.utils.data.DataLoader(trainset, batch_size=14, shuffle=True, generator=_course_generator)
testloader = torch.utils.data.DataLoader(testset, batch_size=14, shuffle=False)
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

display_dataset(trainset,classes=classes)

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.conv3 = nn.Conv2d(16,120,5)
        self.flat = nn.Flatten()
        self.fc1 = nn.Linear(120,64)
        self.fc2 = nn.Linear(64,10)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = nn.functional.relu(self.conv3(x))
        x = self.flat(x)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = LeNet()

summary(net,input_size=(1,3,32,32))

opt = torch.optim.SGD(net.parameters(),lr=0.001,momentum=0.9)
hist = train(net, trainloader, testloader, epochs=3, optimizer=opt, loss_fn=nn.CrossEntropyLoss())