# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ### In this Notebook I develop a first model on MNIST dataset using plain PyTorch

# +
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torch.utils.data import random_split

import torch.optim as optim
from torchvision.datasets import MNIST
from torchvision import transforms

import matplotlib.pyplot as plt
# -

# globals
BATCH_SIZE = 64

# +
# data
dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())

mnist_train, mnist_val = random_split(dataset, [55000, 5000])

train_loader = DataLoader(mnist_train, batch_size=BATCH_SIZE)
val_loader = DataLoader(mnist_val, batch_size=BATCH_SIZE)

# +
# have a look at inputs

for i, batch in enumerate(train_loader):
    inputs, targets = batch
    
    print(inputs.shape)
    
    if i == 0:
        break

# +
# now let's define a simple network to train on MNIST
N_INPUT = 28*28*1
N_CLASSES = 10

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(N_INPUT, 84)
        self.fc2 = nn.Linear(84, 50)
        self.fc3 = nn.Linear(50,N_CLASSES)
    
    def forward(self, x):
        x = x.view(-1, N_INPUT)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # network produces Logits
        x = self.fc3(x)
        
        return x


# +
simplenet = SimpleNet()

optimizer = optim.Adam(simplenet.parameters(), lr=0.001)

# +
# custom train loop

history = []

def train(model, optimizer, loss_fn, train_loader, val_loader, epochs=5, device="cpu"):
    for epoch in range(1, epochs+1):
        training_loss = 0.0
        valid_loss = 0.0
        
        model.train()
        
        for batch in train_loader:
            optimizer.zero_grad()
            inputs, targets = batch
            
            # portiamo dati sulla GPU, se disponibile
            inputs = inputs.to(device)
            targets = targets.to(device)
            output = model(inputs)
            
            loss = loss_fn(output, targets)
            
            # back propagation
            loss.backward()
            optimizer.step()
            
            training_loss += loss.data.item() * inputs.size(0)
        # compute the average
        training_loss /= len(train_loader.dataset)
        
        model.eval()
        
        num_correct = 0 
        num_examples = 0
        
        for batch in val_loader:
            inputs, targets = batch
            inputs = inputs.to(device)
            
            output = model(inputs)
            
            targets = targets.to(device)
            
            loss = loss_fn(output,targets) 
            valid_loss += loss.data.item() * inputs.size(0)
            
            # compute accuracy on validation set
            correct = torch.eq(torch.max(F.softmax(output, dim=1), dim=1)[1], targets)
            num_correct += torch.sum(correct).item()
            num_examples += correct.shape[0]
        valid_loss /= len(val_loader.dataset)

        accuracy = num_correct / num_examples
        
        print('Epoch: {}, Training Loss: {:.2f}, Validation Loss: {:.2f}, accuracy = {:.2f}'.format(epoch, training_loss,
                                                                                                    valid_loss, accuracy))
        # for history during epochs
        metrics = {}
        metrics['loss'] = training_loss
        metrics['valid_loss'] = valid_loss
        metrics['accuracy'] = accuracy
        
        history.append(metrics)
        
    return history


# -

history = train(simplenet, optimizer,torch.nn.CrossEntropyLoss(), train_loader, val_loader, epochs=10, device='cpu')

history

# +
# extract the data from the history list
vet_loss = [x['loss'] for x in history]
vet_val_loss = [x['valid_loss'] for x in history]

plt.figure(figsize=(9,6))
plt.title('Loss')
plt.xlabel('epochs')
plt.ylabel('loss')
plt.plot(vet_loss, label='training loss')
plt.plot(vet_val_loss, label='validation loss')
plt.legend()
plt.grid()

# +
# extract the data from the history list
vet_acc = [x['accuracy'] for x in history]

plt.figure(figsize=(9,6))
plt.title('Accuracy')
plt.xlabel('epochs')
plt.ylabel('loss')
plt.plot(vet_acc, label='validation accuracy')
plt.legend()
plt.grid()
# -



StarcoderdataJupyterScriptsDedupFiltered