%load_ext autoreload
%autoreload 2
%matplotlib inline

import torch
from torch import tensor
from torch import nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from pathlib import Path
from IPython.core.debugger import set_trace
from fastai import datasets
import pickle, gzip, math, torch
import operator

MNIST_URL='http://deeplearning.net/data/mnist/mnist.pkl'

def get_data():
    path = datasets.download_data(MNIST_URL, ext='.gz')
    with gzip.open(path, 'rb') as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')
    return map(tensor, (x_train, y_train, x_valid, y_valid))

x_train, y_train, x_valid, y_valid = get_data()

# x = data, m = mean, s = standard deviation
def normalize(x, m, s): 
    return (x-m)/s

train_mean, train_std = x_train.mean(), x_train.std()
x_train = normalize(x_train, train_mean, train_std)
x_valid = normalize(x_valid, train_mean, train_std)

n, m = x_train.shape
c = y_train.max()+1
n, m, c

class Model(nn.Module):
    # number of input feature, number of hidden feature, number of output feature
    def __init__(self, n_in, nh, n_out):
        super().__init__()
#         set_trace()
        self.layers = [nn.Linear(n_in, nh), nn.ReLU(), nn.Linear(nh, n_out)]        
    def __call__(self, x):
        for l in self.layers:
            x = l(x)
        return x

# loss function
loss_func = F.cross_entropy

def accuracy(yhat, y):
    return (torch.argmax(yhat, dim=-1) == y).float().mean()

# batch size
bs = 64
# learning rate
lr = 0.03
epoch = 1
nh = 50

# x batch
xb = x_train[:bs]
yb = y_train[:bs]

model = Model(m, nh, c.numpy())

# predictions
preds = model(xb)
preds[0]

preds[0].shape

loss = loss_func(preds, yb)
loss

accuracy(preds, yb)

losses, metrics = [], []

def get_databatch(b):
    start_b = b*bs
    end_b = start_b+bs
    return x_train[start_b:end_b], y_train[start_b:end_b]

# e = epoch number
for e in range(epoch):
    # b = batch number
    for b in range((n-1)//bs + 1):
        
        # Feedforward
        xb, yb = get_databatch(b)
        yhatb = model(xb)
        loss = loss_func(yhatb, yb)

        # Metrics
        acc = accuracy(yhatb, yb)
        losses.append(loss); metrics.append(acc)
        
        # Backpropagation
        loss.backward()
        
        # Update Wight and Bias with SGD
        with torch.no_grad():
            for l in model.layers:
                if hasattr(l, 'weight'):
                    l.weight -= l.weight.grad * lr
                    l.bias   -= l.bias.grad * lr 
                    l.weight.grad.zero_()
                    l.bias.grad.zero_()       
                    

x = torch.arange(len(losses)).numpy()
fig,ax = plt.subplots(figsize=(9, 9))
ax.grid(True)
ax.plot(x, losses, label="Loss")
ax.plot(x, metrics, label="Accuracy")
ax.legend(loc='upper right')

loss_func(model(xb), yb), accuracy(model(xb), yb)

xb = x_valid[:bs*2]
yb = y_valid[:bs*2]

yhat = model(xb)
loss_func(yhat, yb), accuracy(yhat, yb)

plt.imshow(xb[99].view(28,28))

yhat[99].argmax()

plt.imshow(xb[55].view(28,28))

yhat[55].argmax()