import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import torchvision
import torchvision.transforms as transforms

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

batch_size = 50

train_dataset = torchvision.datasets.MNIST(root='./data',
                                           train=True,
                                           transform=transforms.ToTensor(),
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root='./data',
                                          train=False,
                                          transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

examples = iter(test_loader)
example_data, example_targets = next(examples)

for i in range(6):
    plt.subplot(2,3,i+1)
    plt.imshow(example_data[i][0], cmap='gray')
    print(example_targets[i].item(),end = " ")
plt.show()

7 2 1 0 4 1

input_size = 784 # 28x28
hidden_size = 200
num_classes = 10
epochs = 2
learning_rate = 0.01
# losses = torch.zeros(epochs*len(train_loader)//10)

# Model = nn.Sequential(
#     nn.Linear(input_size,hidden_size,device=device),
#     nn.ReLU(),
#     nn.Linear(hidden_size, num_classes, device=device),
# )

class ANN_Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(input_size,hidden_size,device=device)
        self.linear2 = nn.Linear(hidden_size, num_classes, device=device)

    def forward(self, x):
        x = self.linear1(x)
        x = F.relu(x)
        x = self.linear2(x)
        return x

Model = ANN_Model().to(device)

# Loss = nn.MSELoss()
# Loss = nn.BCELoss()
Loss = nn.CrossEntropyLoss () # applies nn.LogSoftmax + nn.NLLLoss, No softmax in last layer

# init optimizer
# optim = torch.optim.SGD(Model.parameters(), lr=learning_rate,momentum=0.9)
optim = torch.optim.Adam(Model.parameters(), lr=learning_rate)

n_total_steps = len(train_loader)

for epoch in range(epochs):
    for i, (images, lables) in enumerate(train_loader):

        # origin shape: [50, 1, 28, 28]
        # resized: [50, 784]
        images = images.reshape(-1, 28*28).to(device)
        lables = lables.to(device)

        # forward
        Y_predict = Model(images)

        # backward
        loss = Loss(Y_predict, lables)
        loss.backward()

        # update wights
        optim.step()

        if i % 100 == 0:
            print(f"epoch= {epoch} i= {i+epoch * n_total_steps} loss={loss.item():.4f} ")

        # zero grads
        optim.zero_grad()

epoch= 0 i= 0 loss=2.3142 
epoch= 0 i= 100 loss=0.3428 
...
epoch= 1 i= 2300 loss=0.2180

with torch.no_grad():
    n_correct = 0
    n_samples = 0
    for images, lables in test_loader:
        images = images.reshape(-1, 28*28).to(device)
        lables = lables.to(device)
        __,y_predict = torch.max(Model(images),1)
        n_samples += lables.size(0)
        n_correct += (y_predict == lables).sum().item()

    acc = 100 * n_correct / n_samples
    print(f'Accuracy of the network on the {n_samples} test images: {acc} %')

Accuracy of the network on the 10000 test images: 96.01 %

examples = iter(test_loader)
example_data, example_targets = next(examples)
example_data = example_data.to(device)
example_targets = example_targets.to(device)
_, example_predict_arg = torch.max(Model(example_data.reshape(-1, 28*28)),1)
# print(example_predict_arg)
for i in range(len(example_data)):
    print ("Predict: ",example_predict_arg[i].item(), "target Label: ",example_targets[i].item() )

Predict:  7 target Label:  7
Predict:  2 target Label:  2
Predict:  1 target Label:  1
Predict:  0 target Label:  0
Predict:  4 target Label:  4
Predict:  1 target Label:  1
Predict:  4 target Label:  4
Predict:  9 target Label:  9
Predict:  6 target Label:  5
Predict:  9 target Label:  9
Predict:  0 target Label:  0
Predict:  6 target Label:  6
Predict:  9 target Label:  9
Predict:  0 target Label:  0
Predict:  1 target Label:  1
Predict:  5 target Label:  5
Predict:  9 target Label:  9
Predict:  7 target Label:  7
Predict:  8 target Label:  3
Predict:  4 target Label:  4
Predict:  9 target Label:  9
Predict:  6 target Label:  6
Predict:  6 target Label:  6
Predict:  5 target Label:  5
Predict:  4 target Label:  4
Predict:  0 target Label:  0
Predict:  7 target Label:  7
Predict:  4 target Label:  4
Predict:  0 target Label:  0
Predict:  1 target Label:  1
Predict:  3 target Label:  3
Predict:  1 target Label:  1
Predict:  3 target Label:  3
Predict:  4 target Label:  4
Predict:  7 target Label:  7
Predict:  2 target Label:  2
Predict:  7 target Label:  7
Predict:  1 target Label:  1
Predict:  2 target Label:  2
Predict:  1 target Label:  1
Predict:  1 target Label:  1
Predict:  7 target Label:  7
Predict:  4 target Label:  4
Predict:  2 target Label:  2
Predict:  3 target Label:  3
Predict:  5 target Label:  5
Predict:  1 target Label:  1
Predict:  2 target Label:  2
Predict:  4 target Label:  4
Predict:  4 target Label:  4

File = "MINST_Model.pth"
torch.save (Model, File)

new_model = torch.load(File, weights_only=False)
new_model.eval()

examples = iter(test_loader)
example_data, example_targets = next(examples)
example_data = example_data.to(device)
example_targets = example_targets.to(device)
_, example_predict_arg = torch.max(new_model(example_data.reshape(-1, 28*28)),1)
# print(example_predict_arg)
for i in range(len(example_data)):
    print ("Predict: ",example_predict_arg[i].item(), "target Label: ",example_targets[i].item() )

FILE2 = "model2.pth"
torch.save(Model.state_dict(), FILE2)

new_model2 = ANN_Model().to(device)
state_dict = torch.load(FILE2)
new_model2.load_state_dict(state_dict)
new_model2.eval()

examples = iter(test_loader)
example_data, example_targets = next(examples)
example_data = example_data.to(device)
example_targets = example_targets.to(device)
_, example_predict_arg = torch.max(new_model2(example_data.reshape(-1, 28*28)),1)
# print(example_predict_arg)
for i in range(len(example_data)):
    print ("Predict: ",example_predict_arg[i].item(), "target Label: ",example_targets[i].item() )

FILE3 = "/content/drive/MyDrive/Colab Notebooks/Models/model2.pth"
torch.save(Model, FILE3)