Transfer Learning Py Torch

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

# 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

from pathlib import Path

course_data_dir = Path("data/course-pets")

def course_pet_image(kind, index, size=224):
    rng = np.random.default_rng(10_000 + index + (0 if kind == "Cat" else 1_000))
    y, x = np.mgrid[0:size, 0:size]
    image = np.full((size, size, 3), 0.12, dtype="float32")
    image += rng.normal(0, 0.025, image.shape).astype("float32")
    cx, cy = size // 2 + rng.integers(-12, 13, size=2)
    face = (x - cx) ** 2 + (y - cy) ** 2 < (size * 0.27) ** 2
    if kind == "Cat":
        ears = ((y < cy - size * 0.18) & (np.abs(x - cx) > size * 0.12) & (np.abs(x - cx) < size * 0.3))
        shape = face | ears
        color = np.array([0.85, 0.58, 0.25], dtype="float32")
    else:
        ears = (((x - (cx - size * 0.3)) / (size * 0.14)) ** 2 + ((y - cy) / (size * 0.25)) ** 2 < 1) | (((x - (cx + size * 0.3)) / (size * 0.14)) ** 2 + ((y - cy) / (size * 0.25)) ** 2 < 1)
        shape = face | ears
        color = np.array([0.35, 0.62, 0.88], dtype="float32")
    image[shape] = color
    return np.clip(image, 0, 1)

for class_name in ("Cat", "Dog"):
    class_dir = course_data_dir / class_name
    class_dir.mkdir(parents=True, exist_ok=True)
    for index in range(48):
        pixels = (course_pet_image(class_name, index) * 255).astype("uint8")
        Image.fromarray(pixels).save(class_dir / f"{index:03}.png")

data_dir = str(course_data_dir)

print(f"Prepared the deterministic course fixture at {data_dir}.")

check_image_dir(f"{data_dir}/Cat/*.png")
check_image_dir(f"{data_dir}/Dog/*.png")

std_normalize = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225],
)
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    std_normalize,
])
dataset = torchvision.datasets.ImageFolder(data_dir, transform=transform)
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
split_generator = torch.Generator().manual_seed(2026)
trainset, testset = torch.utils.data.random_split(
    dataset, [train_size, test_size], generator=split_generator
)

display_dataset(dataset, classes=dataset.classes)

# Keep the release notebook deterministic and offline: pretrained weights
# are intentionally not fetched by the embedded execution profile.
vgg = torchvision.models.vgg16(weights=None)
sample_image = dataset[0][0].unsqueeze(0)
res = vgg(sample_image)
print(res[0].argmax())

class_map = {index: (f"course-{index:04d}", f"fixture class {index}") for index in range(1000)}
class_map[res[0].argmax().item()]

summary(vgg,input_size=(1,3,224,224))

device = 'cuda' if torch.cuda.is_available() else 'cpu'

print('Doing computations on device = {}'.format(device))

vgg.to(device)
sample_image = sample_image.to(device)

vgg(sample_image).argmax()

res = vgg.features(sample_image).cpu()
plt.figure(figsize=(15,3))
plt.imshow(res.detach().view(512,-1).T)
print(res.size())

bs = 8
dl = torch.utils.data.DataLoader(dataset,batch_size=bs,shuffle=True)
num = len(dataset)
feature_tensor = torch.zeros(num,512*7*7).to(device)
label_tensor = torch.zeros(num).to(device)
i = 0
for x,l in dl:
    with torch.no_grad():
        f = vgg.features(x.to(device))
        batch_count = len(l)
        feature_tensor[i:i+batch_count] = f.view(batch_count, -1)
        label_tensor[i:i+batch_count] = l
        i += batch_count
        print('.',end='')
        if i>=num:
            break

vgg_dataset = torch.utils.data.TensorDataset(feature_tensor,label_tensor.to(torch.long))
train_count = max(1, int(0.8 * len(vgg_dataset)))
test_count = len(vgg_dataset) - train_count
train_ds, test_ds = torch.utils.data.random_split(vgg_dataset, [train_count, test_count], generator=torch.Generator().manual_seed(2026))

train_loader = torch.utils.data.DataLoader(train_ds,batch_size=32)
test_loader = torch.utils.data.DataLoader(test_ds,batch_size=32)

net = torch.nn.Sequential(torch.nn.Linear(512*7*7,2),torch.nn.LogSoftmax(dim=1)).to(device)

history = train(net,train_loader,test_loader)

print(vgg)

vgg.classifier = torch.nn.Linear(25088,2).to(device)

for x in vgg.features.parameters():
    x.requires_grad = False

summary(vgg,(1, 3,244,244))

trainset, testset = torch.utils.data.random_split(dataset, [max(1, int(0.8 * len(dataset))), len(dataset) - max(1, int(0.8 * len(dataset)))], generator=torch.Generator().manual_seed(42))
train_loader = torch.utils.data.DataLoader(trainset,batch_size=16)
test_loader = torch.utils.data.DataLoader(testset,batch_size=16)

train_long(vgg,train_loader,test_loader,loss_fn=torch.nn.CrossEntropyLoss(),epochs=1)

torch.save(vgg.state_dict(), 'data/cats_dogs.pth')

vgg.load_state_dict(torch.load('data/cats_dogs.pth', map_location=device, weights_only=True))
vgg = vgg.to(device)

for x in vgg.features.parameters():
    x.requires_grad = True

train_long(vgg,train_loader,test_loader,loss_fn=torch.nn.CrossEntropyLoss(),epochs=1,lr=0.0001)

resnet = torchvision.models.resnet18()
print(resnet)