# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.7.3 64-bit (''.env'': conda)'
#     language: python
#     name: python3
# ---

# +
import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import numpy as np
import pandas as pd
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
from tqdm import tqdm
# -

transform = transforms.ToTensor()

train_data = datasets.MNIST('./data/', train=True, download=True, transform=transform)

test_data = datasets.MNIST('./data/', train=False, download=True, transform=transform)

train_data

test_data

# +
img, label = train_data[0]

print(img.shape)
print(label)
# -

plt.imshow(img[0,:,:], cmap='gist_yarg')

# +
torch.manual_seed(101)

train_loader = DataLoader(train_data, batch_size=100, shuffle=True)
test_loader = DataLoader(test_data, batch_size=100)

# +
from torchvision.utils import make_grid

for images, labels in train_loader:
    break

images.shape, labels.shape


# -

# ## Model

class MultilayerPerceptron(nn.Module):
    def __init__(self, in_sz=28*28, out_sz=10, layers=[120, 84]):
        super().__init__()

        self.fc1 = nn.Linear(in_sz, layers[0])
        self.fc2 = nn.Linear(layers[0], layers[1])
        self.out = nn.Linear(layers[1], out_sz)

    def forward(self, X):
        X = F.relu(self.fc1(X))
        X = F.relu(self.fc2(X))
        X = self.out(X)

        return F.softmax(X, dim=1)


# +
torch.manual_seed(101)

model = MultilayerPerceptron()
model
# -

print('in:', 28*28)
for param in model.parameters():
    print(param.numel())
print('-> outs')
    # 784 pixels -> number of connections -> number of layer nodes (with no bias) -> etc... -> 10 is outs

# Above gives 105214 parameters in the model!

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# flattening images from [num_batches, 1, 28, 28] to [num_batches, 784]
images.shape, images.view(100, -1).shape

# +
epochs = 10
train_losses = []
train_correct = []
test_losses = []
test_correct = []

for i in range(epochs):

    train_corr = 0
    test_corr = 0

    for batch, (X_train, y_train) in enumerate(train_loader):

        batch += 1

        y_pred = model(X_train.view(100, -1))

        loss = criterion(y_pred, y_train)

        predicted_class = torch.max(y_pred.data, 1)[1]
        batch_accuracy = (predicted_class == y_train).sum()
        train_corr += batch_accuracy

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        if batch % 200 == 0:
            acc = train_corr.item() * 100 / (100 * batch)
            print(f'TRAIN Epoch: {i}, batch: {batch}, loss: {loss.item()}, accuracy: {acc}')

    train_losses.append(loss)
    train_correct.append(train_corr)

    with torch.no_grad():
        for batch, (X_test, y_test) in enumerate(test_loader):

            y_val = model(X_test.view(100, -1))

            predicted_class_test = torch.max(y_val.data, 1)[1]
            test_corr += (predicted_class_test == y_test).sum()

    test_loss = criterion(y_val, y_test)
    test_losses.append(test_loss)
    test_correct.append(test_corr)
# -

plt.plot([i.detach().numpy() for i in train_losses], label='train loss')
plt.plot([i.detach().numpy() for i in test_losses], label='test loss')
plt.legend()

train_acc = [t/600 for t in train_correct]
test_acc = [t/100 for t in test_correct]

plt.plot(train_acc, label='train acc')
plt.plot(test_acc, label='test acc')
plt.legend()



StarcoderdataJupyterScriptsDedupFiltered