pip install jovian --upgrade

import numpy as np
import torch
from torchvision.datasets import MNIST

dataset = MNIST(root = '/data', download = True)

len(dataset)

test_data = MNIST(root = '/data', train = False)
len(test_data)

import matplotlib.pyplot as plt
%matplotlib inline

image, label = dataset[0]
plt.imshow(image, cmap = 'gray')
print('Label:', label)

image, label = dataset[10]
plt.imshow(image, cmap = 'gray')
print('Label:', label)

#To convert images into tensors
import torchvision.transforms as transforms

#Dataset after converting into tensors
dataset = MNIST(root = '/data',
               train = True,
               transform = transforms.ToTensor())

img_tensor, label = dataset[0]
print(img_tensor.shape, label)

print(img_tensor[:, 10:15, 10:15])
print(torch.max(img_tensor), torch.min(img_tensor))

plt.imshow(img_tensor[0, 10:15, 10:15], cmap = 'gray');

plt.imshow(img_tensor[0], cmap = 'gray');

import numpy as np


def split_indices(n, val_pct):
    
    '''
    Shuffles the array indices 0, 1, ... , n-1 and seperates out a random portion for
    the validation set. np.random.permutation() is used for this purpose.
    '''
    
    #Determine size of validation set
    n_val = int(val_pct * n)
    
    #Create random permutations of 0 to n-1
    idxs = np.random.permutation(n)
    
    #Pick first n_val for validation set
    return idxs[n_val:], idxs[:n_val]

train_idxs, validation_idxs = split_indices(len(dataset), 0.2)

print(len(train_idxs), len(validation_idxs))

#Create PyTorch data loaders for each of the training and validation indices
#USe SubsetRandomSampler(), which samples elements randomly from a given list of indices
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data.dataloader import DataLoader

batch_size = 100

#Train sampler and data loader
train_sampler = SubsetRandomSampler(train_idxs)
train_loader = DataLoader(dataset, 
                         batch_size,
                         sampler = train_sampler)

#Validation sampler and data loader
val_sampler = SubsetRandomSampler(validation_idxs)
val_loader = DataLoader(dataset,
                       batch_size,
                       sampler = val_sampler)

import torch.nn as nn

input_size = 28*28
num_classes = 10

class MnistModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(input_size, num_classes)

    def forward(self, xb):
        xb = xb.reshape(-1, 784)
        out = self.linear(xb)
        return out
    
model = MnistModel()

print(model.linear.weight.shape, model.linear.bias.shape)
list(model.parameters())

for images, labels in train_loader:
    outputs = model(images)
    break
    
print('Output shape:', outputs.shape)
print('Sample outputs:\n', outputs[:2].data)

#Applying the activation function to convert the logit scores to probabilities
import torch.nn.functional as F

#Apply softmax to each output row
probs = F.softmax(outputs, dim=1)

print('Sample probabilities:', probs[:2].data)
print('Sum of probabilities:', torch.sum(probs[0]).item())

max_probs, preds = torch.max(probs, dim=1)
print(preds)
print(max_probs)

labels

def accuracy(preds, target):
    return torch.sum(preds == target).item() / len(preds)

accuracy(preds, label)

loss_fn = F.cross_entropy

#Loss for the current batch of data
loss = loss_fn(outputs, labels)
print(loss)

lr = 0.0001
optimizer = torch.optim.Adam(model.parameters(), lr = lr)

def loss_batch(model, loss_fn, xb, yb, opt = None, metric = None):
    #Calculate loss
    preds = model(xb)
    loss = loss_fn(preds, yb)
    
    if opt is not None:
        #Compute gradients
        loss.backward()
        
        #Update parameters
        opt.step()
        
        #Reset gradients
        opt.zero_grad()
        
    metric_res = None
    if metric is not None:
        metric_res = metric(preds, yb)
        
    return loss.item(), len(xb), metric_res

def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.sum(preds == labels).item() / len(preds)

def evaluate(model, loss_fn, valid_dl, metric = None):
    with torch.no_grad():
        #Pass each batch through the model
        results = [loss_batch(model, loss_fn, xb, yb, metric=accuracy) for xb, yb in valid_dl]
        
        #Seperate losses, counts and metrics
        loss, nums, metrics = zip(*results)
        
        #Total size of the dataset
        total = np.sum(nums)
        
        #Avg loss across all batches
        avg_loss = np.sum(np.multiply(loss, nums)) / total
        
        avg_metric = None
        if metric is not None:
            avg_metric = np.sum(np.multiply(metrics, nums)) / total
        
    return avg_loss, total, avg_metric

val_loss, total, val_metric = evaluate(model, loss_fn, val_loader, metric = accuracy)
print(val_loss, val_metric)

def fit(epochs, model, loss_fn, opt, train_dl, val_dl, metric = None):
    for epoch in range(epochs):
        #Training
        for xb, yb in train_dl:
            loss,_,_ = loss_batch(model, loss_fn, xb, yb, opt = optimizer, metric = accuracy)
            
        #Evaluation
        result = evaluate(model, loss_fn, val_dl, metric)
        val_loss, total, val_metric = result
        
        #Print progress of training
        if metric is None:
            print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, epochs, val_loss))
        else:
            print('Epoch [{}/{}], Loss: {:.4f}, Accuracy: {:.2f}'.format(epoch+1, epochs, val_loss, val_metric))

#Redefine model and optimizer
model = MnistModel()
optimizer = torch.optim.Adam(model.parameters(), lr = lr)

fit(10, model, F.cross_entropy, optimizer, train_loader, val_loader, metric = accuracy)

acc = [0.89, 0.89, 0.90, 0.90, 0.90, 0.90, 0.90, 0.90, 0.91, 0.91]
plt.plot(acc, '-x')
plt.xlabel('epoch')
plt.ylabel('accuracies')
plt.title('Accuracy vs no of epochs')
plt.show()

!pip install jovian --upgrade

lr

import jovian

jovian.log_hyperparams({
    'opt':'Adam',
    'lr':'0.0001',
    'batch_size':100,
    'arch':'logistic-regression',
    'epochs':10
})

jovian.log_metrics({
    'val_loss':0.3402,
    'val_acc':0.91
})

#Define test dataset
test = MNIST(root = '/data', train = False, transform=transforms.ToTensor())

img, label = test[5]
plt.imshow(img[0], cmap='gray')
print('Shape:', img.shape)
print('Label:', label)

img.unsqueeze(0).shape

def predict_image(img, model):
    xb = img.unsqueeze(0)
    yb = model(xb)
    _, preds = torch.max(yb, dim=1)
    return preds[0].item()

img, label = test[10]
plt.imshow(img[0], cmap='gray')
print('Label:', label, 'Prediction:', predict_image(img, model))

img, label = test[190]
plt.imshow(img[0], cmap='gray')
print('Label:', label, 'Prediction:', predict_image(img, model))

img, label = test[297]
plt.imshow(img[0], cmap='gray')
print('Label:', label, 'Prediction:', predict_image(img, model))

img, label = test[9998]
plt.imshow(img[0], cmap='gray')
print('Label:', label, 'Prediction:', predict_image(img, model))

test_loader = DataLoader(test, batch_size=200)

test_loss, total, test_metric = evaluate(model, F.cross_entropy, test_loader, metric = accuracy)
print('Loss: {:.4f}, Accuracy: {:.2f}'.format(test_loss, test_metric))

jovian.log_metrics({
    'test_loss':0.3206,
    'test_acc':0.91
})

torch.save(model.state_dict(), 'mnist-logistic-adam.pth')

model.state_dict()

model2 = MnistModel()
model2.load_state_dict(torch.load('mnist-logistic-adam.pth'))
model2.state_dict()

test_loss, total, test_metric = evaluate(model2, F.cross_entropy, test_loader, metric = accuracy)
print('Loss: {:.4f}, Accuracy: {:.2f}'.format(test_loss, test_metric))

jovian.commit(artifacts=['mnist-logistic-adam.pth'])