Neural networks 2

# this call is done only when working with .ipynb
%load_ext tensorboard

from torch.utils.tensorboard import SummaryWriter

# default `log_dir` is "runs" - we'll be more specific here
writer = SummaryWriter('logs')

# this call is done only when working with .ipynb
%tensorboard --logdir logs

for i in range(100):
    writer.add_scalar("y=2x", i * 2, i)

# imports
import matplotlib.pyplot as plt
import numpy as np

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

# transforms
transform = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))])

# datasets
trainset = torchvision.datasets.FashionMNIST('./data',
    download=True,
    train=True,
    transform=transform)
testset = torchvision.datasets.FashionMNIST('./data',
    download=True,
    train=False,
    transform=transform)

# dataloaders
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                        shuffle=True, num_workers=2)


testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                        shuffle=False, num_workers=2)

# constant for classes
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
        'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')

print(f" is cuda available: {torch.cuda.is_available()}")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

fig,ax = plt.subplots(6,10,figsize=(10,10))
ax = ax.flatten()

for i in range(len(ax)):
    ax[i].axis('off')
    ax[i].imshow(trainset.data[i,:,:], cmap='gray')
    ax[i].set_title(classes[int(trainset.targets[i])])
plt.tight_layout()

trainset.data.shape

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        # TRY IT: explain yourself the inputs and outputs sizes

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()
net.to(device)
net

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

def eval_network(net):
    net.eval()

    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in testloader:
            # Move batch to device
            images = images.to(device)
            labels = labels.to(device)

            outputs = net(images)
            predicted = torch.argmax(outputs.data, dim=1)
            total += labels.shape[0]
            correct += (predicted == labels).sum()
    
    accuracy_prc = 100 * correct / total

    # set back to train mode
    net.train()
    
    return accuracy_prc

accuracy_prc = eval_network(net)
print(f"Accuracy of the network test images: {accuracy_prc}%")

def plot_test_results_to_figure(net, num_images=4):
    # Ensure the network is in evaluation mode
    net.eval()

    # Get one batch of images and labels
    images, labels = next(iter(testloader))

    # Only select a subset of images (up to num_images)
    images = images[:num_images]
    labels = labels[:num_images]

    # Move images to the same device as the model
    images = images.to(device)

    # Forward pass to get outputs
    output = net(images)

    # Convert outputs to predicted class indices and probabilities
    _, preds = torch.max(output, 1)
    probs = F.softmax(output, dim=1)

    # Initialize the matplotlib figure
    fig = plt.figure(figsize=(num_images * 3, 3))  # Adjust the size as needed

    # Loop over the images in the batch
    for idx in range(num_images):
        ax = fig.add_subplot(1, num_images, idx + 1, xticks=[], yticks=[])
        img = images[idx].cpu().numpy()  # Convert image to numpy array
        img = np.transpose(img, (1, 2, 0))  # Reorder dimensions for displaying

        # Display the image
        ax.imshow(img, cmap="Greys")

        # Set the title for each image
        true_label = classes[labels[idx].item()]
        pred_label = classes[preds[idx].item()]
        probability = probs[idx][preds[idx]].item() * 100
        ax.set_title(
            f"{pred_label}, {probability:.1f}%\n(GT: {true_label})",
            color=("green" if preds[idx] == labels[idx] else "red"),
        )

    net.train()
    return fig


writer.add_figure(
    "test plot",
    plot_test_results_to_figure(net),
    0,
)

running_loss = 0.0
for epoch in range(1):  # loop over the dataset multiple times
    for i, data in enumerate(trainloader):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # Move batch to device
        inputs = inputs.to(device)
        labels = labels.to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if i % 1000==0:  # every x mini-batches...
            print(f"{i}/{len(trainloader)}")
            global_step = epoch * len(trainloader) + i
            
            # ...log the running loss
            writer.add_scalar(
                "training loss", running_loss / 1000, global_step
            )
            running_loss = 0.0

            # ... log test accuracy
            accuracy_prc = eval_network(net)
            writer.add_scalar(
                "test accuracy", accuracy_prc, global_step
            )

            # ...log a Matplotlib Figure showing the model's predictions on a
            # random mini-batch
            writer.add_figure(
                "prediction vs. GT",
                plot_test_results_to_figure(net),
                global_step,
            )
print("Finished Training")