"""
# Deep Learning in Biomedical Engineering
"""
"""
## Assignment 2
### Due date:<font color='red'> 11:59 pm, Febraury 25, 2021</font>
"""
"""
## Xin Gao xg2326
"""
"""
## Introduction
In this assignment, you will implement, train, and test [LeNet](https://en.wikipedia.org/wiki/LeNet) model to classify [CIFAR10](http://www.cs.toronto.edu/~kriz/cifar.html) dataset.

## Instructions
Depending on each question, there are empty blocks you need to fill. The code blocks are only for Python codes and comments and currently have a comment like ```# [Your code here]```. The markdown blocks are only for plain or laTex text that you may use for answering descriptive questions. Currently, the markdown blocks have a comment like <font color='red'>[your answer here]</font>. Please remove the comments before filling the blocks. You can always add more blocks if you need to, or it just makes your answers well-organized.

Although you may use other available online resources (such as GitHub, Kaggle Notebooks), it is highly recommended to try your best to do it yourself. If you are using an online code or paper, make sure you cite their work properly to avoid plagiarism. Using other students' works is absolutely forbidden and considered cheating.

Write comments for your codes in the big picture mode. One comment line at the top of each code block is necessary to explain what you did in that block. Don't comment on every detail.

Name your variables properly that represent the data they are holding, such as ``` test_set = ..., learning_rate = ...``` not ```a = ..., b = ...```.

Implementing and reporting results using other architectures than LeNet will grant you an extra 20% on grade.

In this [Kaggle Notebook](https://www.kaggle.com/roblexnana/cifar10-with-cnn-for-beginer), you may find useful information about how your outputs must look. Just remember, they are using a different architecture, and they are using TensorFlow for implementations.


## How to submit:
After you have completed the assignment: 
1. Save a version (You can use 'Quick Save' to avoid re-running the whole notebook)
2. Make the saved version public (Share -> Public)
3. Copy the 'Public URL'
4. Download your completed notebook as a '.ipynb' file (File -> Download)
5. Upload the 'Public URL' and the '.ipynb' files on the [CourseWorks](https://courseworks2.columbia.edu/).
"""
"""
### 1. (5 pts.) According to the [CIFAR10](http://www.cs.toronto.edu/~kriz/cifar.html) dataset descriptions and other online resources, please identify the quantities below:

a) Total Number of samples \
b) Number of classes \
c) Image size \
d) Write class names and their corresponding label index (e.g., 0. cats, 1. dogs, ...) \
e) Intensity range
"""
"""

a) total number of samples: 60,000  
b) number of classes: 10  
c) image size: 3*32*32  
d) 0. airplane 1. automobile 2. bird 3. cat 4. deer 5. dog 6. frog 7. horse 8. ship 9. truck   
e) intensity range: 0-255     

"""
"""
### 2. (0 pts.) Import the required packages in the following code block. 
You can always come back here and import another package; please keep them all in the following block to make your code well-organized.
"""
# import packages
import numpy as np
from matplotlib import pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import torchvision
import torchvision.transforms as transforms

import seaborn as sns
from sklearn.metrics import confusion_matrix, accuracy_score

# for printing training progress
from IPython.display import clear_output
# initialize the seed and define batch size: 
torch.manual_seed(2326)
np.random.seed(2326)

BATCH_SIZE = 32
EPOCHS = 200
# check if GPU is available
cuda = torch.cuda.is_available()
print("GPU available:", cuda)
"""
### 3. (5 pts.) Load train and test sets using Pytorch datasets functions.
Make sure the intensity range is between $[0, 1]$ and images are stored as 'tensor' type. 
You may use transformers while downloading the dataset to do the job. Look at this tutorial: [Pytorch CIFAR10](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html))
"""
# code imported and modified from pytorch tutorial
# load data into trainset and testset

transform = transforms.ToTensor()

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)

# project numerical values to actual classes

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

NUM_ROUNDS = len(trainset) / BATCH_SIZE
!mkdir LeNet_MSE
"""
### 4. (10 pts.) Using the 'matplotlib' library, make a figure with $N\times4$ grid cells where $N$ is the number of classes. Display one random sample of each class pulled from the train set in its corresponding row that depends on its class index and the first column and its histogram in the second column. Repeat this for the third and fourth columns but pull images from the test set.
"""
# show histogram fuction

def getHistograms(img, intensity_range,num_bins=10):
    histograms = list()
    for i in range(3):
        hist, edge = np.histogram(img[:,:,i], bins=num_bins, range=intensity_range)
        edge = (edge[:-1] + edge[1:]) /2 # use the middle point as bin representation
        histograms.append((edge, hist))
    return histograms
# Display one random sample of each class and its histogram

# Because each batch only has 32 samples, it's possible that each batch won't have
# any image of some classes

# Define trainset data loader
dataiter = iter(torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
                                          shuffle=True, num_workers=1))

# load the first batch of images and find the first image with label==0
images, labels = dataiter.next()
index = (np.where(labels == 0))

plt.figure(figsize=(20, 60))

#reference https://thispointer.com/find-the-index-of-a-value-in-numpy-array/

for i in range(10): # for each class
    
    # find until batch has images in this category, store in images and labels
    while len(np.where(labels == i)[0]) == 0:
        images, labels = dataiter.next()
    
    # select the first image in this class to be displayed
    index = np.where(labels == i)[0]
    image = images[index][0,:,:,:]
    
    #display the image
    plt.subplot(10,4,4*i+1)
    plt.title(classes[i]+' - traning set')
    npimg = np.transpose(image.numpy(), (1, 2, 0)) 
    plt.imshow(npimg)
    
    # show the histogram in each channel
    plt.subplot(10,4,4*i+2)
    plt.title('Histogram')
    histograms = getHistograms(npimg, intensity_range=(0.0, 1.0), num_bins=10)
    colors = ('Red', 'Green', 'Blue')
    for i in range(3):
        plt.plot(histograms[i][0],histograms[i][1],c=colors[i], alpha=0.8,label='Channel {0}'.format(colors[i]))
    plt.xlabel('Intensity')
    plt.ylabel('Abundance')
    plt.legend()
    
    
    
# repeat the process with the testset
dataiter = iter(torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE,
                                          shuffle=True, num_workers=1))
images, labels = dataiter.next()
index = (np.where(labels == 0))

for i in range(10):
    
    while len(np.where(labels == i)[0]) == 0:
        images, labels = dataiter.next()

    index = np.where(labels == i)[0]
    image = images[index][0,:,:,:]
    
    plt.subplot(10,4,4*i+3)
    plt.title(classes[i]+' - test set')
    npimg = np.transpose(image.numpy(), (1, 2, 0)) 
    plt.imshow(npimg)
    
    plt.subplot(10,4,4*i+4)
    plt.title('Histogram')
    histograms = getHistograms(npimg, intensity_range=(0.0, 1.0), num_bins=10)
    colors = ('Red', 'Green', 'Blue')
    for i in range(3):
        plt.plot(histograms[i][0],histograms[i][1],c=colors[i], alpha=0.8,label='Channel {0}'.format(colors[i]))
    plt.xlabel('Intensity')
    plt.ylabel('Abundance')
    plt.legend()
        

plt.show() 
    
    
"""
### 5. (5 pts.) Split up the train set into new train and validation sets so that the number of samples in the validation set is equal to the number of samples in the test set. Then create a [```DataLoader```](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) for each set, including the test set, with a batch size of 32.
The [Pytorch CIFAR10](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html) tutorial will also help you here.
Make sure none of the samples in the validation set exists in the new train set.
"""
#https://medium.com/@sergioalves94/deep-learning-in-pytorch-with-cifar-10-dataset-858b504a6b54
# validation set size equal to test set size

vali_size = len(testset) 

# rest of the train set become the new train set

train_size = len(trainset) - vali_size 
trainset, valiset = torch.utils.data.random_split(trainset, [train_size, vali_size])

#

trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
                                          shuffle=True, num_workers=1)
valiloader = torch.utils.data.DataLoader(valiset, batch_size=BATCH_SIZE,
                                          shuffle=True, num_workers=1)
testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE,
                                         shuffle=False, num_workers=1)
"""
### 6. (5 pts.) Display the number of samples for each class in the train, validation, and test sets as a stacked bar plot similar to the '[FirstTutorial](https://www.kaggle.com/soroush361/deeplearninginbme-firsttutorial)'.
"""
# given a data loader, find the number of samples in each classes in the whole dataset

def countLoader(dataloader):
    
    count = np.zeros(10)
    
    for i, data in enumerate(dataloader):
    
        inputs, labels = data
                    
        for j in range(labels.shape[0]):
            count[labels[j]] += 1
        
    return count
print("No. of samples in: \n Train set: {} \n Validation set: {} \n Test set: {}".format(len(trainset), len(valiset),len(testset)))

train_count = countLoader(trainloader)

vali_count = countLoader(valiloader)

test_count = countLoader(testloader)
plt.figure(figsize=(12,6))
plt.title('data distribution')
plt.xticks(range(10), labels=range(10))
plt.xlabel('Number')
plt.ylabel('Number of samples')
plt.bar(range(10), train_count, label='Train')
plt.bar(range(10), vali_count, label='Validation', bottom=train_count)
plt.bar(range(10), test_count, label='Test', bottom=vali_count+train_count)
plt.legend()
plt.savefig('DataDistribution.png')
plt.show()
"""
### 7. (10 pts.) According to the LeNet architecture below, create a fully-connected model. Also, identify the architeture's hyper-parameters, activation functions, and tensor shapes.
Architecture hyper-parameter includes the number of layers, number of kernels in each layer, size of the kernels, stride, zero-padding size. Just by looking at the architecture itself, you should be able to identify the hyper-parameters. Keep in mind that you identified the $W, H$, and $N$ (Which refers to the number of classes) in the first question.
For more help, look at this [implementation](https://github.com/icpm/pytorch-cifar10/blob/master/models/LeNet.py).
![LeNet Architecture](https://raw.githubusercontent.com/soroush361/DeepLearningInBME/main/Ass1_Arch1.png)
"""
# Formula to calculate shape as we go through layer by layer = [(X - F + 2P)/S] + 1
# Here,
# X = Width / Height
# F = Kernel size
# P = Padding
# S = Strides (default = 1)

# Our input to the first layer is going to be [batchsize, 1, 32, 32]
# substitute, =[(32 - 5 + 2(0))/1] + 1
#             = 28

#reference: https://github.com/icpm/pytorch-cifar10/blob/master/models/LeNet.py

class Model(nn.Module):
    '''ConvNet -> Max_Pool -> RELU -> ConvNet -> Max_Pool -> RELU -> FC -> RELU -> FC -> SOFTMAX'''
    def __init__(self):
        '''Define model modules.'''
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x)) # convolution 1
        x = F.max_pool2d(x, 2)    # max pooling 1
        x = F.relu(self.conv2(x)) # cconvolution 2
        x = F.max_pool2d(x, 2)    # max pooling 2
        x = x.view(x.size(0), -1) # flatterning
        x = F.relu(self.fc1(x))   # reshape
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

LeNet = Model()

# If GPU available, move the model to GPU.
if cuda:
    LeNet.cuda()

print(LeNet)
"""
 <font color='red'>[Describe the model's architecture hyper-parameters]</font>
"""
"""
Conv1 + MaxPool1:   
    kernel size = 5*5  
    stride = 1  
    no zero-padding  
    maxpool: 2*2 with stride=2  
    input channel=3  
    output channel=6   
    
Conv2 + MaxPool2:   
    kernel size = 5*5  
    stride = 1  
    no zero-padding  
    maxpool: 2*2 with stride=2  
    input channel=6  
    output channel=16  
"""
"""
### 8. (5 pts.) Create an instance of [ADAM optimizer](https://pytorch.org/docs/stable/optim.html#torch.optim.Adam) with an initial learning rate of $0.0001$ and an instance of [Mean Squared Error (MSE)](https://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html#torch.nn.MSELoss) loss function. Briefly explain the ADAM optimizer algorithm and the MSE loss function.
For ADAM optimizer, keep other arguments as default. 

For your information, here is the mathematics behind the ADAM optimizer: \
For each parameter $w^j$
$$
v_t = \beta_1v_{t-1}-(1-\beta_1)g_t \\
s_t = \beta_2s_{t-1}-(1-\beta_2)g_t^2 \\
\Delta w^j = -\eta\frac{v_t}{\sqrt{s_t+\epsilon}}g_t \\
w^j_{t+1} = w^j_t+\Delta w^j
$$
Where $\eta$ is the initial learning rate, $g_t$ is the gradient at time $t$ along $w^j$, $v_t$ is the exponential average of gradients along $w^j$, $s_t$ is the exponential average of squares of gradients along $w^j$, $\beta_1, \beta_2$ are the hyper-parameters, and $\epsilon$ is a small number to avoid dividing by zero.

The MSE loss function is:
$$
L(y,\hat{y}) = \frac{1}{N}\sum_{i=1}^N{(y_i-\hat{y}_i)^2}
$$
Where $y$ is the true value, $\hat{y}$ is the predicted value, $N$ is the number of classes. Keep in mind that $y$ is a one-hot vector like $y=\begin{bmatrix} 0 \\ 0 \\ 1 \\ \vdots \end{bmatrix}$ (This example of $y$ indicates that the sample belongs to class ID 2, remember it is zero-indexed) and $\hat{y}=\begin{bmatrix} 0.1 \\ 0.03 \\ 0.8 \\ \vdots \end{bmatrix}$ shows the probability of belonging to each class for the same sample and predicted by the model.

For other optimization algorithms and loss functions, check the links below:

[Optimizers list](https://pytorch.org/docs/stable/optim.html#algorithms) 

[Loss function list](https://pytorch.org/docs/stable/nn.html#loss-functions)
"""
# [Create your optimiser and loss function here]
optimizer = optim.Adam(LeNet.parameters(), lr=0.0001)
criterion = nn.MSELoss()
"""
<font color='red'>[Explaine ADAM and MSE here]</font>
"""
"""
Adam is a optimized gradien descent that combines the AdaGrad and RMSProp algorithms to provide an optimization algorithm that can handle sparse gradients on noisy problems.

MSE is just the mean square error loss function. It represents the second moment of the error and is prompt to large varience by small perturbation in certain regions.
"""
"""
### 9. (15 pts.) Train the model for 200 epochs using the train set and validate your model's performance at the end of each epoch using the validation sets.
To increase the training speed, use the GPU accelerator.
Look at the '[FirstTutorial](https://www.kaggle.com/soroush361/deeplearninginbme-firsttutorial)' for more help.
Do not forget to save the model at the end of each epoch.
"""
print('Started Training')

train_epoch_loss = []
validation_epoch_loss = []

for epoch in range(EPOCHS):  # loop over the dataset multiple times

    train_loss = []
    validation_loss = []
    loss = 0.0 # loss is a temporary variable used in loops
    
    print('Progress: {:2.0f}%'.format(epoch/EPOCHS*100),end=' ')
    clear_output(wait=True)
    
    #######################################################
    ####################### Train #########################
    #######################################################
    # Set the model to train mode so that the parameters can be updated.
        
    LeNet.train()
    
    for i, data in enumerate(trainloader):
        
        # If GPU is available, move the data to the GPU for faster computation.
        if cuda:
                
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data
            labels = F.one_hot(labels.cuda(), num_classes=10).float()

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = LeNet(inputs.cuda())
            
            loss = criterion(outputs.cuda(), labels.cuda())
            train_loss.append(loss.cpu().data.item())
            
            loss.backward()
            
            optimizer.step()

                                    
        else:
           
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data
            labels = F.one_hot(labels, num_classes=10).float()

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = LeNet(inputs)
            
            loss = criterion(outputs, labels)
            train_loss.append(loss.cpu().data.item())
            
            loss.backward()
            
            optimizer.step()
            
    #######################################################
    ###################### Validation #####################
    #######################################################
    # Set the model to evaluation mode so that parameters are fixed.          
    
    LeNet.eval()
    
    for j, validation_data in enumerate(valiloader):           
            
        if cuda:

            # get the inputs; data is a list of [inputs_validation, labels_validation]
            validation_input, validation_labels = validation_data
            validation_labels = F.one_hot(validation_labels.cuda(), num_classes=10).float()
                
            validation_output = LeNet(validation_input.cuda())
                
            loss = criterion(validation_output.cuda(), validation_labels.cuda())
            validation_loss.append(loss.cpu().data.item()) # why cpu?
        
        else:
          
            # get the inputs; data is a list of [inputs_validation, labels_validation]
            validation_input, validation_labels = validation_data
            validation_labels = F.one_hot(validation_labels, num_classes=10).float()
                
            validation_output = LeNet(validation_input)
                
            loss = criterion(validation_output, validation_labels)
            validation_loss.append(loss.cpu().data.item())
                
    train_epoch_loss.append(np.mean(train_loss))
    validation_epoch_loss.append(np.mean(validation_loss))
    
   
    # save models
    torch.save(LeNet.state_dict(), './LeNet_MSE/checkpoint_epoch_%s.pth' % (epoch))
    
print('Finished Training')
"""
### 10. (5 pts.) Display the learning curve and illustrate the best epoch. Explain your criteria for the best epoch.
The learning curve shows the model's loss and accuracy at the end of each epoch for all epochs (200 epochs). The criteria for the best epoch can be the minimum loss or maximum accuracy or other criteria.
"""
plt.figure(figsize = (12, 8))
plt.plot(train_epoch_loss, '-o', label = 'training loss', markersize = 3)
plt.plot(validation_epoch_loss, '-o', label = 'validation loss', markersize = 3)
plt.legend(loc = 'upper right');
best_epoch = np.argmin(validation_epoch_loss)
print('best epoch: ', best_epoch)
"""
<font color='red'>[Describe your criteria for choosing the best epoch]</font>
"""
"""
the lowest validation error to allows for best generalization results
"""
"""
### 11. (10 pts.) Load the model's weights at the best epoch and test your model performance on the test set. Display the confusion matrix and classification report.
"""
# [Your code here for loading the best model, testing on the test set, and displaying the confusion matrix as well as the classification report]
# load the best model
state_dict = torch.load('./LeNet_MSE/checkpoint_epoch_%s.pth' % (best_epoch))
print(state_dict.keys())
LeNet.load_state_dict(state_dict)
full_test_labels = []
full_predicted_labels = []
full_images = []
        
LeNet.eval()

for i, test_data in enumerate(testloader):
    
    if cuda:
    
        # get the inputs; data is a list of [test_inputs, test_labels]
        test_input, test_labels = test_data
        test_labels = F.one_hot(test_labels.cuda(), num_classes=10).float()
                
        test_output = LeNet(test_input.cuda())
        
        for j in range(test_labels.shape[0]):
            full_test_labels.append(test_labels.tolist()[j])
            full_predicted_labels.append(test_output.tolist()[j])
            full_images.append(test_input.tolist()[j])
                
       
    else:

        # get the inputs; data is a list of [test_inputs, test_labels]
        test_input, test_labels = test_data
        test_labels = F.one_hot(test_labels, num_classes=10).float()
                
        test_output = LeNet(test_input)
        
        for j in range(test_labels.shape[0]):
            full_test_labels.append(test_labels.tolist()[j])
            full_predicted_labels.append(test_output.tolist()[j])
            full_images.append(test_input.tolist()[j])
# check number of samples tested
print("number of samples tested:",len(full_test_labels))

# convert probability into prediction labels

full_test_labels_num = np.argmax(full_test_labels, axis=1)
full_predicted_labels_num = np.argmax(full_predicted_labels, axis=1)

# an array of whether the prediction of each element is right

test_result = full_test_labels_num == full_predicted_labels_num


# manual confusion matrix output

confusion_matrix_num = np.zeros((10,10))
for i in range(len(full_test_labels)):
    
    confusion_matrix_num[full_test_labels_num[i], full_predicted_labels_num[i]] += 1
    
print(confusion_matrix_num)


confusion_matrix_num.trace()/ len(testset)
print("Total correct rate is: {:2.2f} %".format(confusion_matrix_num.trace()/len(testset)*100))
CM = confusion_matrix(full_test_labels_num, full_predicted_labels_num)

plt.figure(figsize = (12,10))
sns.heatmap(CM, annot = True, annot_kws = {"size": 10})
plt.ylim([0, 10]);
plt.ylabel('True labels');
plt.xlabel('predicted labels');
"""
### 12. (5 pts.) Display five random samples of each class titled with the true label and the predicted label. Comment on your model's performance.
Samples must be pulled from the test set.
"""
# [Your code here]
#reshape the data for visualization

print("The original dimension of the list is: {} * {} * {} * {}".format(len(full_images),len(full_images[0]),len(full_images[0][0]), len(full_images[0][0][0])))

full_images_swaped = np.swapaxes(np.swapaxes(full_images,1,3),1,2)

print("The dimension of the list after transform is: {} * {} * {} * {}".format(len(full_images_swaped),len(full_images_swaped[0]),len(full_images_swaped[0][0]), len(full_images_swaped[0][0][0])))
# generate one image from each class with labels and predictions

plt.figure(figsize=(20, 60))
for i in range(10):
    class_index = np.where(full_test_labels_num == i)                  # find the images with label i
    
    random_indexes = np.random.choice(class_index[0], replace=True, size=5)   # randomly choose 5 samples from the list
    for j in range(5):
        
        random_index = random_indexes[j]
        plt.subplot(10, 5, i*5+j+1)
        plt.title("Actual Label:{} \n Predicted Label:{}".format(classes[full_test_labels_num[random_index]],
                                                            classes[full_predicted_labels_num[random_index]]))
        plt.imshow(full_images_swaped[random_index])
"""
<font color='red'>[Comment about the model's performance]</font>
"""
"""
This model achieved ~60% error rate with 2 epochs, 50% error rate with 10 epochs
"""
"""
### 13. (20 pts.) Repeat the training, validation, and testing with the [Cross-Entropy](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss) loss function and initial learning rate of $0.005$. Explain how the model's performance changed.
Essentially you need to copy all the codes above and just change the loss function and edit the learning rate. Obviously, you don't need to re-import the dataset and the libraries. However, you need to create a new instance of the architecture. Otherwise, the weights would be the same as the last epoch (or the best epoch) in the last part. To avoid overwriting your previously trained model, change the save directories in the training loop.
"""
# [Your code here]
!mkdir LeNet_CrossEntropy
print("No. of samples in: \n Train set: {} \n Validation set: {} \n Test set: {}".format(len(trainset), len(valiset),len(testset)))
#reference: https://github.com/icpm/pytorch-cifar10/blob/master/models/LeNet.py

class Model2(nn.Module):
    '''ConvNet -> Max_Pool -> RELU -> ConvNet -> Max_Pool -> RELU -> FC -> RELU -> FC -> SOFTMAX'''
    def __init__(self):
        '''Define model modules.'''
        super(Model2, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x)) # convolution 1
        x = F.max_pool2d(x, 2)    # max pooling 1
        x = F.relu(self.conv2(x)) # cconvolution 2
        x = F.max_pool2d(x, 2)    # max pooling 2
        x = x.view(x.size(0), -1) # flatterning
        x = F.relu(self.fc1(x))   # reshape
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

LeNet2 = Model2()

# If GPU available, move the model to GPU.
if cuda:
    LeNet2.cuda()

print(LeNet2)
optimizer = optim.Adam(LeNet2.parameters(), lr=0.005)
criterion = nn.CrossEntropyLoss()
print('Started Training')

train_epoch_loss = []
validation_epoch_loss = []

for epoch in range(EPOCHS):  # loop over the dataset multiple times

    train_loss = []
    validation_loss = []
    loss = 0.0 # loss is a temporary variable used in loops
    
    print('Progress: {:2.0f}%'.format(epoch/EPOCHS*100),end=' ')
    clear_output(wait=True)
    
    #######################################################
    ####################### Train #########################
    #######################################################
    # Set the model to train mode so that the parameters can be updated.
        
    LeNet2.train()
    
    for i, data in enumerate(trainloader):
        
        # If GPU is available, move the data to the GPU for faster computation.
        if cuda:
                
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = LeNet2(inputs.cuda())
            
            loss = criterion(outputs.cuda(), labels.cuda())
            train_loss.append(loss.cpu().data.item())
            
            loss.backward()
            
            optimizer.step()

                                    
        else:
           
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = LeNet2(inputs)
            
            loss = criterion(outputs, labels)
            train_loss.append(loss.cpu().data.item())
            
            loss.backward()
            
            optimizer.step()
            
    #######################################################
    ###################### Validation #####################
    #######################################################
    # Set the model to evaluation mode so that parameters are fixed.          
    
    LeNet2.eval()
    
    for j, validation_data in enumerate(valiloader):           
            
        if cuda:

            # get the inputs; data is a list of [inputs_validation, labels_validation]
            validation_input, validation_labels = validation_data
                
            validation_output = LeNet2(validation_input.cuda())
                
            loss = criterion(validation_output.cuda(), validation_labels.cuda())
            validation_loss.append(loss.cpu().data.item()) # why cpu?
        
        else:
          
            # get the inputs; data is a list of [inputs_validation, labels_validation]
            validation_input, validation_labels = validation_data
                
            validation_output = LeNet2(validation_input)
                
            loss = criterion(validation_output, validation_labels)
            validation_loss.append(loss.cpu().data.item())
                
    train_epoch_loss.append(np.mean(train_loss))
    validation_epoch_loss.append(np.mean(validation_loss))
    
   
    # save models
    torch.save(LeNet2.state_dict(), './LeNet_CrossEntropy/checkpoint_epoch_%s.pth' % (epoch))
    
print('Finished Training')
plt.figure(figsize = (12, 8))
plt.plot(train_epoch_loss, '-o', label = 'training loss', markersize = 3)
plt.plot(validation_epoch_loss, '-o', label = 'validation loss', markersize = 3)
plt.legend(loc = 'upper right');
best_epoch = np.argmin(validation_epoch_loss)
print('best epoch: ', best_epoch)

state_dict = torch.load('./LeNet_CrossEntropy/checkpoint_epoch_%s.pth' % (best_epoch))
print(state_dict.keys())
LeNet2.load_state_dict(state_dict)
full_test_labels = []
full_predicted_labels = []
full_images = []

LeNet2.eval()

for i, test_data in enumerate(testloader):
    
    if cuda:

        # get the inputs; data is a list of [test_inputs, test_labels]
        test_input, test_labels = test_data
        test_labels = F.one_hot(test_labels.cuda(), num_classes=10).float()
                
        test_output = LeNet2(test_input.cuda())
        
        for j in range(test_labels.shape[0]):
            full_test_labels.append(test_labels.tolist()[j])
            full_predicted_labels.append(test_output.tolist()[j])
            full_images.append(test_input.tolist()[j])
                
       
    else:

        # get the inputs; data is a list of [test_inputs, test_labels]
        test_input, test_labels = test_data
        test_labels = F.one_hot(test_labels, num_classes=10).float()
                
        test_output = LeNet2(test_input)
        
        for j in range(test_labels.shape[0]):
            full_test_labels.append(test_labels.tolist()[j])
            full_predicted_labels.append(test_output.tolist()[j])
            full_images.append(test_input.tolist()[j])
print("Number of samples tested: {}".format(len(full_test_labels)))

# convert probability into final label 

full_test_labels_num = np.argmax(full_test_labels, axis=1)
full_predicted_labels_num = np.argmax(full_predicted_labels, axis=1)

# an array of whether the prediction of each element is right
test_result = full_test_labels_num == full_predicted_labels_num

# manual output the confusion matrix 

confusion_matrix_num = np.zeros((10,10))
for i in range(len(full_test_labels)):
    
    confusion_matrix_num[full_test_labels_num[i], full_predicted_labels_num[i]] += 1
    
print(confusion_matrix_num)
confusion_matrix_num.trace()/ len(testset)
print("Total correct rate is: {:2.2f} %".format(confusion_matrix_num.trace()/len(testset)*100))
CM = confusion_matrix(full_test_labels_num, full_predicted_labels_num)

plt.figure(figsize = (12,10))
sns.heatmap(CM, annot = True, annot_kws = {"size": 10})
plt.ylim([0, 10]);
plt.ylabel('True labels');
plt.xlabel('predicted labels');
print("The original dimension of the list is: {} * {} * {} * {}".format(len(full_images),len(full_images[0]),len(full_images[0][0]), len(full_images[0][0][0])))

#reshape the data for visualization
full_images_swaped = np.swapaxes(np.swapaxes(full_images,1,3),1,2)

print("The dimension of the list after transform is: {} * {} * {} * {}".format(len(full_images_swaped),len(full_images_swaped[0]),len(full_images_swaped[0][0]), len(full_images_swaped[0][0][0])))
# generate one image from each class with labels and predictions

plt.figure(figsize=(20, 60))
for i in range(10):
    class_index = np.where(full_test_labels_num == i)                  # find the images with label i
    
    random_indexes = np.random.choice(class_index[0], replace=True, size=5)   # randomly choose 5 samples from the list
    for j in range(5):
        
        random_index = random_indexes[j]
        plt.subplot(10, 5, i*5+j+1)
        plt.title("Actual Label:{} \n Predicted Label:{}".format(classes[full_test_labels_num[random_index]],
                                                            classes[full_predicted_labels_num[random_index]]))
        plt.imshow(full_images_swaped[random_index])
"""
<font color='red'>[Comment about the performance changes]</font>
"""
"""
This model achieved ~62% error rate with 2 epochs, ~51% error rate with 10 epochs.

However after 200 epochs the model based on cross entropy failed to converge. This probably indicates that the learning rate is too high and model doesn't generalize well compared to the previous model.
"""