%pip install torch torchvision

import torch
import torchvision
import numpy as np
from torchvision import datasets, transforms
from torch.utils.data.sampler import SubsetRandomSampler
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import os
import torch.quantization
import json
from PIL import Image
import copy
import time
from torch.optim import lr_scheduler

N, D = 14, 10
x = torch.randn(N, D).type(torch.FloatTensor)
print(x)

from torchvision import models

alexnet = models.alexnet()
print(alexnet)

# check if CUDA is available
train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
    print("CUDA is not available.  Training on CPU ...")
else:
    print("CUDA is available!  Training on GPU ...")

# number of subprocesses to use for data loading
num_workers = 0
# how many samples per batch to load
batch_size = 20
# percentage of training set to use as validation
valid_size = 0.2

# convert data to a normalized torch.FloatTensor
transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)

# choose the training and test datasets
train_data = datasets.CIFAR10("data", train=True, download=True, transform=transform)
test_data = datasets.CIFAR10("data", train=False, download=True, transform=transform)

# obtain training indices that will be used for validation
num_train = len(train_data)
indices = list(range(num_train))
np.random.shuffle(indices)
split = int(np.floor(valid_size * num_train))
train_idx, valid_idx = indices[split:], indices[:split]

# define samplers for obtaining training and validation batches
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetRandomSampler(valid_idx)

# prepare data loaders (combine dataset and sampler)
train_loader = torch.utils.data.DataLoader(
    train_data, batch_size=batch_size, sampler=train_sampler, num_workers=num_workers
)
valid_loader = torch.utils.data.DataLoader(
    train_data, batch_size=batch_size, sampler=valid_sampler, num_workers=num_workers
)
test_loader = torch.utils.data.DataLoader(
    test_data, batch_size=batch_size, num_workers=num_workers
)

# specify the image classes
classes = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
]

# define the CNN architecture


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 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 = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# create a complete CNN
model_1 = Net()
print(model_1)
# move tensors to GPU if CUDA is available
if train_on_gpu:
    model_1.cuda()

criterion = nn.CrossEntropyLoss()  # specify loss function
optimizer = optim.SGD(model_1.parameters(), lr=0.01)  # specify optimizer

n_epochs = 30  # number of epochs to train the model
train_loss_list = []  # list to store loss to visualize
valid_loss_list = []
valid_loss_min = np.Inf  # track change in validation loss

for epoch in range(n_epochs):
    # Keep track of training and validation loss
    train_loss = 0.0
    valid_loss = 0.0

    # Train the model
    model_1.train()
    for data, target in train_loader:
        # Move tensors to GPU if CUDA is available
        if train_on_gpu:
            data, target = data.cuda(), target.cuda()
        # Clear the gradients of all optimized variables
        optimizer.zero_grad()
        # Forward pass: compute predicted outputs by passing inputs to the model
        output = model_1(data)
        # Calculate the batch loss
        loss = criterion(output, target)
        # Backward pass: compute gradient of the loss with respect to model parameters
        loss.backward()
        # Perform a single optimization step (parameter update)
        optimizer.step()
        # Update training loss
        train_loss += loss.item() * data.size(0)

    # Validate the model
    model_1.eval()
    for data, target in valid_loader:
        # Move tensors to GPU if CUDA is available
        if train_on_gpu:
            data, target = data.cuda(), target.cuda()
        # Forward pass: compute predicted outputs by passing inputs to the model
        output = model_1(data)
        # Calculate the batch loss
        loss = criterion(output, target)
        # Update average validation loss
        valid_loss += loss.item() * data.size(0)

    # Calculate average losses
    train_loss = train_loss / len(train_loader)
    valid_loss = valid_loss / len(valid_loader)
    train_loss_list.append(train_loss)
    valid_loss_list.append(valid_loss)

    # Print training/validation statistics
    print(
        "Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}".format(
            epoch, train_loss, valid_loss
        )
    )

    # Save model if validation loss has decreased
    if valid_loss <= valid_loss_min:
        print(
            "Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...".format(
                valid_loss_min, valid_loss
            )
        )
        torch.save(model_1.state_dict(), "model_1_cifar.pt")
        valid_loss_min = valid_loss

fig_1,ax_1 = plt.subplots()
ax_1.plot(range(n_epochs), train_loss_list, label = 'train_loss')
ax_1.legend()
ax_1.plot(range(n_epochs), valid_loss_list, label = 'valid_loss')
ax_1.legend()
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Performance of Model 1")
plt.show()

 model_1.load_state_dict(torch.load("./model_1_cifar.pt"))

# track test loss
test_loss = 0.0
class_correct = list(0.0 for i in range(10))
class_total = list(0.0 for i in range(10))
accuracy_per_class_m1 =  list(0.0 for i in range(10))
model_1.eval()
# iterate over test data
for data, target in test_loader:
    # move tensors to GPU if CUDA is available
    if train_on_gpu:
        data, target = data.cuda(), target.cuda()
    # forward pass: compute predicted outputs by passing inputs to the model
    output = model_1(data)
    # calculate the batch loss
    loss = criterion(output, target)
    # update test loss
    test_loss += loss.item() * data.size(0)
    # convert output probabilities to predicted class
    _, pred = torch.max(output, 1)
    # compare predictions to true label
    correct_tensor = pred.eq(target.data.view_as(pred))
    correct = (
        np.squeeze(correct_tensor.numpy())
        if not train_on_gpu
        else np.squeeze(correct_tensor.cpu().numpy())
    )
    # calculate test accuracy for each object class
    for i in range(batch_size):
        label = target.data[i]
        class_correct[label] += correct[i].item()
        class_total[label] += 1

# average test loss
test_loss = test_loss / len(test_loader)
print("Test Loss: {:.6f}\n".format(test_loss))

for i in range(10):
    if class_total[i] > 0:
        accuracy_per_class_m1 = 100 * class_correct[i] / class_total[i]
        print(
            "Test Accuracy of %5s: %2d%% (%2d/%2d)"
            % (
                classes[i],
                accuracy_per_class_m1,
                np.sum(class_correct[i]),
                np.sum(class_total[i]),
            )
        )
        
    else:
        print("Test Accuracy of %5s: N/A (no training examples)" % (classes[i]))

print(
    "\nTest Accuracy of model 1 (Overall): %2f%% (%2f/%2f)"
    % (
        100.0 * np.sum(class_correct) / np.sum(class_total),
        np.sum(class_correct),
        np.sum(class_total),
    )
)

# define the CNN architecture


class Second_Net(nn.Module):
    def __init__(self):
        super(Second_Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        # H_out(conv1) = W_out(conv1) = (H_in + 2*padding - Kernel)/stride + 1 = (32 + 2 - 3)/1 + 1 = 32
        # Size(conv1_out) = 32*32*16
        self.pool = nn.MaxPool2d(2, 2)  # After Maxpooling2D new size of image equal =  16*16
        # H_out(conv2) = W_out(conv2) = (H_in + 2*padding - Kernel)/stride + 1 = (16 + 2 - 3)/1 + 1 = 16
                                        # After Maxpooling2D new size of image equal =  8*8
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        # H_out(conv3) = W_out(conv3) = (H_in + 2*padding - Kernel)/stride + 1 = (16 + 2 - 3)/1 + 1 = 8
                                        # After Maxpooling2D new size of image equal =  4*4
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)  
        self.fc1 = nn.Linear(64 * 4 * 4, 512)   
        self.fc2 = nn.Linear(512, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 64 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# create a complete CNN
model_2 = Second_Net()
print(model_2)
# move tensors to GPU if CUDA is available
# if train_on_gpu:
  #  model.cuda()

criterion = nn.CrossEntropyLoss()  # specify loss function
optimizer = optim.SGD(model_2.parameters(), lr=0.01)  # specify optimizer
train_loss_list_2 = []  # list to store loss to visualize
valid_loss_list_2 = []
valid_loss_min = np.Inf  # track change in validation loss

for epoch in range(n_epochs):
    # Keep track of training and validation loss
    train_loss = 0.0
    valid_loss = 0.0

    # Train the model
    model_2.train()
    for data, target in train_loader:
        # Move tensors to GPU if CUDA is available
        if train_on_gpu:
            data, target = data.cuda(), target.cuda()
        # Clear the gradients of all optimized variables
        optimizer.zero_grad()
        # Forward pass: compute predicted outputs by passing inputs to the model
        output = model_2(data)
        # Calculate the batch loss
        loss = criterion(output, target)
        # Backward pass: compute gradient of the loss with respect to model parameters
        loss.backward()
        # Perform a single optimization step (parameter update)
        optimizer.step()
        # Update training loss
        train_loss += loss.item() * data.size(0)

    # Validate the model
    model_2.eval()
    for data, target in valid_loader:
        # Move tensors to GPU if CUDA is available
        if train_on_gpu:
            data, target = data.cuda(), target.cuda()
        # Forward pass: compute predicted outputs by passing inputs to the model
        output = model_2(data)
        # Calculate the batch loss
        loss = criterion(output, target)
        # Update average validation loss
        valid_loss += loss.item() * data.size(0)

    # Calculate average losses
    train_loss = train_loss / len(train_loader)
    valid_loss = valid_loss / len(valid_loader)
    train_loss_list_2.append(train_loss)
    valid_loss_list_2.append(valid_loss)

    # Print training/validation statistics
    print(
        "Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}".format(
            epoch, train_loss, valid_loss
        )
    )

    # Save model if validation loss has decreased
    if valid_loss <= valid_loss_min:
        print(
            "Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...".format(
                valid_loss_min, valid_loss
            )
        )
        torch.save(model_2.state_dict(), "model_2_cifar.pt")
        valid_loss_min = valid_loss

fig_2,ax_2 = plt.subplots()
ax_2.plot(range(n_epochs), train_loss_list_2, label = 'train_loss')
ax_2.legend()
ax_2.plot(range(n_epochs), valid_loss_list_2, label = 'valid_loss')
ax_2.legend()
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Performance of Model 2")
plt.show()

model_2.load_state_dict(torch.load("./model_2_cifar.pt"))

class_correct2 = list(0.0 for i in range(10))
class_total2 = list(0.0 for i in range(10))
accuracy_per_class_m2 = list(0.0 for i in range(10))
model_2.eval()
# iterate over test data
for data, target in test_loader:
    # move tensors to GPU if CUDA is available
    if train_on_gpu:
        data, target = data.cuda(), target.cuda()
    # forward pass: compute predicted outputs by passing inputs to the model
    output = model_2(data)

    # convert output probabilities to predicted class
    _, pred = torch.max(output, 1)
    # compare predictions to true label
    correct_tensor = pred.eq(target.data.view_as(pred))
    correct2 = (
        np.squeeze(correct_tensor.numpy())
        if not train_on_gpu
        else np.squeeze(correct_tensor.cpu().numpy())
    )
    # calculate test accuracy for each object class
    for i in range(batch_size):
        label = target.data[i]
        class_correct2[label] += correct2[i].item()
        class_total2[label] += 1


for i in range(10):
    if class_total2[i] > 0:
        accuracy_per_class_m2 =  100 * class_correct2[i] / class_total2[i]
        print(
            "Test Accuracy of %5s: %2d%% (%2d/%2d)"
            % (
                classes[i],
                accuracy_per_class_m2,
                np.sum(class_correct2[i]),
                np.sum(class_total2[i]),
            )
        )
       
    else:
        print("Test Accuracy of %5s: N/A (no training examples)" % (classes[i]))

print(
    "\nTest Accuracy (Overall): %2f%% (%2f/%2f)"
    % (
        100.0 * np.sum(class_correct2) / np.sum(class_total2),
        np.sum(class_correct2),
        np.sum(class_total2),
    )
)

def print_size_of_model(model, label=""):
    torch.save(model.state_dict(), "temp.p")
    size = os.path.getsize("temp.p")
    print("model: ", label, " \t", "Size (KB):", size / 1e3)
    os.remove("temp.p")
    return size


print_size_of_model(model_1, "fp32")

quantized_model_1 = torch.quantization.quantize_dynamic(model_1, dtype=torch.qint8)
print_size_of_model(quantized_model_1, "int8")

print_size_of_model(model_2, "fp32")
quantized_model_2 = torch.quantization.quantize_dynamic(model_2, dtype=torch.qint8)
print_size_of_model(quantized_model_2, "int8")

# quantized_model_1.load_state_dict(torch.load("./quantizedmodel_cifar.pt"))



class_correct = list(0.0 for i in range(10))
class_total = list(0.0 for i in range(10))
accuracy_per_class_q1 = list(0.0 for i in range(10))
quantized_model_1.eval()
# iterate over test data
for data, target in test_loader:
    # move tensors to GPU if CUDA is available
    if train_on_gpu:
        data, target = data.cuda(), target.cuda()
    # forward pass: compute predicted outputs by passing inputs to the model
    output = quantized_model_1(data)
   
    
    # convert output probabilities to predicted class
    _, pred = torch.max(output, 1)
    # compare predictions to true label
    correct_tensor = pred.eq(target.data.view_as(pred))
    correct = (
        np.squeeze(correct_tensor.numpy())
        if not train_on_gpu
        else np.squeeze(correct_tensor.cpu().numpy())
    )
    # calculate test accuracy for each object class
    for i in range(batch_size):
        label = target.data[i]
        class_correct[label] += correct[i].item()
        class_total[label] += 1

for i in range(10):
    if class_total[i] > 0:
        print(
            "Test Accuracy of %5s: %2d%% (%2d/%2d)"
            % (
                classes[i],
                100 * class_correct[i] / class_total[i],
                np.sum(class_correct[i]),
                np.sum(class_total[i]),
            )
        )
        accuracy_per_class_q1[i] = (class_correct[i]/class_total[i])*100
    else:
        print("Test Accuracy of quantized model 1 %5s: N/A (no training examples)" % (classes[i]))

print(
    "\nTest Accuracy of quantized model 1 (Overall): %2f%% (%2f/%2f)"
    % (
        100.0 * np.sum(class_correct) / np.sum(class_total),
        np.sum(class_correct),
        np.sum(class_total),
    )
)

#  quantized_model_2.load_state_dict(torch.load("./model_cifar.pt"))



class_correct = list(0.0 for i in range(10))
class_total = list(0.0 for i in range(10))
accuracy_per_class_q2 = list(0.0 for i in range(10))
quantized_model_2.eval()
# iterate over test data
for data, target in test_loader:
    # move tensors to GPU if CUDA is available
    if train_on_gpu:
        data, target = data.cuda(), target.cuda()
    # forward pass: compute predicted outputs by passing inputs to the model
    output = quantized_model_2(data)
   
    
    # convert output probabilities to predicted class
    _, pred = torch.max(output, 1)
    # compare predictions to true label
    correct_tensor = pred.eq(target.data.view_as(pred))
    correct = (
        np.squeeze(correct_tensor.numpy())
        if not train_on_gpu
        else np.squeeze(correct_tensor.cpu().numpy())
    )
    # calculate test accuracy for each object class
    for i in range(batch_size):
        label = target.data[i]
        class_correct[label] += correct[i].item()
        class_total[label] += 1

for i in range(10):
    if class_total[i] > 0:
        print(
            "Test Accuracy of %5s: %2d%% (%2d/%2d)"
            % (
                classes[i],
                100 * class_correct[i] / class_total[i],
                np.sum(class_correct[i]),
                np.sum(class_total[i]),
            )
        )
        accuracy_per_class_q2[i] = (class_correct[i]/class_total[i])*100
    else:
        print("Test Accuracy of quantized model 2 %5s: N/A (no training examples)" % (classes[i]))

print(
    "\nTest Accuracy of quantized model 2 (Overall): %2f%% (%2f/%2f)"
    % (
        100.0 * np.sum(class_correct) / np.sum(class_total),
        np.sum(class_correct),
        np.sum(class_total),
    )
)

fig_4, ax_4 = plt.subplots()
w = 0.4
x = ["airpl.", "auto.", "bird", "car", "deer", "dog", "frog", "horse", "ship", "truck"]
bar1 = np.arange(len(x))
bar2 = [i+w for i in bar1]
ax_4.bar(bar1,accuracy_per_class_m1,w,label="Original model")
ax_4.bar(bar2,accuracy_per_class_q1,w,label="quantized model")
ax_4.set_xticks(bar1+w/2,x)
ax_4.legend()
plt.xlabel("Class")
plt.ylabel("Accuracy")
plt.show()

fig_5, ax_5 = plt.subplots()
w = 0.4
x = ["airpl.", "autom.", "bird", "car", "deer", "dog", "frog", "horse", "ship", "truck"]
bar1 = np.arange(len(x))
bar2 = [i+w for i in bar1]
ax_5.bar(bar1,accuracy_per_class_m2,w,label="Original model")
ax_5.bar(bar2,accuracy_per_class_q2,w,label="Quantized_model")
ax_5.set_xticks(bar1+w/2,x)
ax_5.legend()
plt.xlabel("Class")
plt.ylabel("Accuracy")
plt.show()

# Choose an image to pass through the model
test_image = "dog.PNG"

# Configure matplotlib for pretty inline plots
#%matplotlib inline
#%config InlineBackend.figure_format = 'retina'

# Prepare the labels
with open("imagenet-simple-labels.json") as f:
    labels = json.load(f)

# First prepare the transformations: resize the image to what the model was trained on and convert it to a tensor
data_transform = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)
# Load the image

image = Image.open(test_image)
plt.imshow(image), plt.xticks([]), plt.yticks([])

# Now apply the transformation, expand the batch dimension, and send the image to the GPU
# image = data_transform(image).unsqueeze(0).cuda()
image = data_transform(image).unsqueeze(0)

# Download the model if it's not there already. It will take a bit on the first run, after that it's fast
model_3 = models.resnet50(pretrained=True)  #  pretrained=True : permet d'obtenir les valeurs des poids après le pré-entrainement.
print(model_3)
# Send the model to the GPU
# model.cuda()
# Set layers such as dropout and batchnorm in evaluation mode
model_3.eval()

# Get the 1000-dimensional model output
out = model_3(image)
# Find the predicted class
print("Predicted class is: {}".format(labels[out.argmax()]))

print_size_of_model(model_3, "fp32")

quantized_model_3 = torch.quantization.quantize_dynamic(model_3, dtype=torch.qint8)
print_size_of_model(quantized_model_3, "int8")

import json
from PIL import Image

# Choose an image to pass through the model
test_image = "dog.PNG"

# Configure matplotlib for pretty inline plots
#%matplotlib inline
#%config InlineBackend.figure_format = 'retina'

# Prepare the labels
with open("imagenet-simple-labels.json") as f:
    labels = json.load(f)

# First prepare the transformations: resize the image to what the model was trained on and convert it to a tensor
data_transform = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)
# Load the image

image = Image.open(test_image)
plt.imshow(image), plt.xticks([]), plt.yticks([])

# Now apply the transformation, expand the batch dimension, and send the image to the GPU
# image = data_transform(image).unsqueeze(0).cuda()
image = data_transform(image).unsqueeze(0)

# Download the model if it's not there already. It will take a bit on the first run, after that it's fast
model_3 = models.resnet50(pretrained=True)  #  pretrained=True : allows you to obtain the weight values after the pre-training.
quantized_model_3 = torch.quantization.quantize_dynamic(model_3, dtype=torch.qint8)
print_size_of_model(quantized_model_3, "int8")
print(quantized_model_3)
# Send the model to the GPU
# model.cuda()
# Set layers such as dropout and batchnorm in evaluation mode
quantized_model_3.eval()

# Get the 1000-dimensional model output
out = quantized_model_3(image)
# Find the predicted class
print("Predicted class is: {}".format(labels[out.argmax()]))

import json
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import transforms, models

# List of test images:
image_list = ["Plane.jpg", "Moon.jpg", "eagle.jpeg", "Audi_Q8.jpg"]

# Prepare the labels
with open("imagenet-simple-labels.json") as f:
    labels = json.load(f)

# First prepare the transformations: resize the image to what the model was trained on and convert it to a tensor
data_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# Load model
model_3 = models.resnet50(pretrained=True)
model_3.eval()

# Set up subplots
fig, axes = plt.subplots(1, len(image_list), figsize=(15, 3))

# Loop through the image list
for i, image_path in enumerate(image_list):
    # Load image
    image = Image.open(image_path)
    
    # Apply transformations
    image_tensor = data_transform(image).unsqueeze(0)
    
    # Perform inference with the model
    with torch.no_grad():
        out = model_3(image_tensor)
    
    # Display the image with its predicted class
    axes[i].imshow(image)
    axes[i].set_title(f"Predicted class: {labels[out.argmax()]}")
    axes[i].axis("off")

# Adjust layout
plt.tight_layout()
plt.show()    

# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
    "train": transforms.Compose(
        [
            transforms.RandomResizedCrop(
                224
            ),  # ImageNet models were trained on 224x224 images
            transforms.RandomHorizontalFlip(),  # flip horizontally 50% of the time - increases train set variability
            transforms.ToTensor(),  # convert it to a PyTorch tensor
            transforms.Normalize(
                [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
            ),  # ImageNet models expect this norm
        ]
    ),
    "val": transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    
    ),
    "test": transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    
    ),
}

data_dir = "hymenoptera_data"
# Create train and validation datasets and loaders
image_datasets = {
    x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
    for x in ["train", "val", "test"]
}
dataloaders = {
    x: torch.utils.data.DataLoader(
        image_datasets[x], batch_size=4, shuffle=True, num_workers=0
    )
    for x in ["train", "val", "test"]
}
dataset_sizes = {x: len(image_datasets[x]) for x in ["train", "val", "test"]}
class_names = image_datasets["train"].classes
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Helper function for displaying images
def imshow(inp, title=None):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])

    # Un-normalize the images
    inp = std * inp + mean
    # Clip just in case
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated
    plt.show()


# Get a batch of training data
inputs, classes = next(iter(dataloaders["train"]))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])

# Get a batch of training data
# inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
# out = torchvision.utils.make_grid(inputs)

# imshow(out, title=[class_names[x] for x in classes])
# training
train_loss_list_7 = []
valid_loss_list_7 = []
train_accuracy_list_7 = []
valid_accuracy_list_7 = []
test_loss_list_7 = []
test_accuracy_list_7 = []

num_epochs = 10

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    epoch_time = []  # we'll keep track of the time needed for each epoch
    
    

    for epoch in range(num_epochs):
        epoch_start = time.time()
        print("Epoch {}/{}".format(epoch + 1, num_epochs))
        print("-" * 10)

        # Each epoch has a training and validation phase
        for phase in ["train", "val"]:
            if phase == "train":
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # Forward
                # Track history if only in training phase
                with torch.set_grad_enabled(phase == "train"):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

                # Statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
         
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            

            if phase == "train" :
                train_loss_list_7.append(epoch_loss)
                train_accuracy_list_7.append(epoch_acc)
            elif phase == "val":
                valid_loss_list_7.append(epoch_loss)
                valid_accuracy_list_7.append(epoch_acc
                )

            print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc))

            # Deep copy the model
            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        # Add the epoch time
        t_epoch = time.time() - epoch_start
        epoch_time.append(t_epoch)
        print()

    time_elapsed = time.time() - since
    print(
        "Training complete in {:.0f}m {:.0f}s".format(
            time_elapsed // 60, time_elapsed % 60
        )
    )
    print("Best val Acc: {:4f}".format(best_acc))

    # Load best model weights
    model.load_state_dict(best_model_wts)
    return model, epoch_time


def test_model (model, criterion):
    since = time.time()

    running_loss = 0.0
    running_corrects = 0
    model.eval()
    for inputs, labels in dataloaders["test"] :
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
    
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
        
    test_loss = running_loss/dataset_sizes["test"]
    test_accuracy = running_corrects.double()/dataset_sizes["test"]
    test_loss_list_7.append(test_loss)
    test_accuracy_list_7.append(test_accuracy)

    print("test loss is %2f and test accuracy  is %2f" %(test_loss, test_accuracy ))

    return test_loss, test_accuracy
    


# Download a pre-trained ResNet18 model and freeze its weights
model_N1 = torchvision.models.resnet18(pretrained=True)
for param in model_N1.parameters():
    param.requires_grad = False

# Replace the final fully connected layer
# Parameters of newly constructed modules have requires_grad=True by default

num_ftrs = model_N1.fc.in_features
model_N1.fc = nn.Linear(num_ftrs, 2)

# Send the model to the GPU
model_N1 = model_N1.to(device)
# Set the loss function
criterion = nn.CrossEntropyLoss()

# Observe that only the parameters of the final layer are being optimized
optimizer_conv = optim.SGD(model_N1.fc.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
model_N1, epoch_time = train_model(
    model_N1, criterion, optimizer_conv, exp_lr_scheduler, num_epochs )

test_loss, test_accuracy = test_model(model_N1, criterion)

fig_6 = plt.figure() 
fig_6.set_size_inches(10, 4)
plt.subplots_adjust(wspace=0.2, hspace=0.35)
ax_6_1 = plt.subplot(1,2,1) 
ax_6_2 = plt.subplot(1,2,2)
ax_6_1.plot(range(num_epochs), train_loss_list_7, label = 'train_loss')
ax_6_1.legend()
ax_6_1.plot(range(num_epochs), valid_loss_list_7, label = 'valid_loss')
ax_6_1.legend()
ax_6_1.set_xlabel("Epoch")
ax_6_1.set_ylabel("Loss")
ax_6_1.set_title("Performance of resnet18 model")
ax_6_2.plot(range(num_epochs), train_accuracy_list_7, label = 'train_acc')
ax_6_2.legend()
ax_6_2.plot(range(num_epochs), valid_accuracy_list_7, label = 'valid_acc')
ax_6_2.legend()
ax_6_2.set_xlabel("Epoch")
ax_6_2.set_ylabel("accuracy")
ax_6_2.set_title("Performance of resnet18 model")
plt.show()

train_loss_list_8 = []
valid_loss_list_8 = []
train_accuracy_list_8 = []
valid_accuracy_list_8 = []
test_loss_list_8 = []
test_accuracy_list_8 = []
hidden_dim = 256
num_epochs = 10

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    epoch_time = []  # we'll keep track of the time needed for each epoch
    
    

    for epoch in range(num_epochs):
        epoch_start = time.time()
        print("Epoch {}/{}".format(epoch + 1, num_epochs))
        print("-" * 10)

        # Each epoch has a training and validation phase
        for phase in ["train", "val"]:
            if phase == "train":
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # Forward
                # Track history if only in training phase
                with torch.set_grad_enabled(phase == "train"):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

                # Statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
         
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            

            if phase == "train" :
                train_loss_list_8.append(epoch_loss)
                train_accuracy_list_8.append(epoch_acc)
            elif phase == "val":
                valid_loss_list_8.append(epoch_loss)
                valid_accuracy_list_8.append(epoch_acc
                )

            print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc))

            # Deep copy the model
            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        # Add the epoch time
        t_epoch = time.time() - epoch_start
        epoch_time.append(t_epoch)
        print()

    time_elapsed = time.time() - since
    print(
        "Training complete in {:.0f}m {:.0f}s".format(
            time_elapsed // 60, time_elapsed % 60
        )
    )
    print("Best val Acc: {:4f}".format(best_acc))

    # Load best model weights
    model.load_state_dict(best_model_wts)
    return model, epoch_time


def test_model (model, criterion):
    since = time.time()

    running_loss = 0.0
    running_corrects = 0
    model.eval()
    for inputs, labels in dataloaders["test"] :
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
    
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
        
    test_loss = running_loss/dataset_sizes["test"]
    test_accuracy = running_corrects.double()/dataset_sizes["test"]
    test_loss_list_8.append(test_loss)
    test_accuracy_list_8.append(test_accuracy)

    print("test loss is %2f and test accuracy  is %2f" %(test_loss, test_accuracy ))

    return test_loss, test_accuracy
    


# Download a pre-trained ResNet18 model and freeze its weights
model_N2 = torchvision.models.resnet18(pretrained=True)
for param in model_N2.parameters():
    param.requires_grad = False

# Replace the final fully connected layer
# Parameters of newly constructed modules have requires_grad=True by default

num_ftrs = model_N2.fc.in_features
model_N2.fc = nn.Sequential(
nn.Dropout(0.2),     
nn.Linear(num_ftrs, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 2))

# Send the model to the GPU
model_N2 = model_N2.to(device)
# Set the loss function
criterion = nn.CrossEntropyLoss()

# Observe that only the parameters of the final layer are being optimized
optimizer_conv = optim.SGD(model_N2.fc.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
model_N2, epoch_time = train_model(
    model_N2, criterion, optimizer_conv, exp_lr_scheduler, num_epochs )

test_loss, test_accuracy = test_model(model_N2, criterion)

fig_6 = plt.figure() 
fig_6.set_size_inches(10, 4)
plt.subplots_adjust(wspace=0.2, hspace=0.35)
ax_6_1 = plt.subplot(1,2,1) 
ax_6_2 = plt.subplot(1,2,2)
ax_6_1.plot(range(num_epochs), train_loss_list_8, label = 'train_loss')
ax_6_1.legend()
ax_6_1.plot(range(num_epochs), valid_loss_list_8, label = 'valid_loss')
ax_6_1.legend()
ax_6_1.set_xlabel("Epoch")
ax_6_1.set_ylabel("Loss")
# ax_6_1.set_title("Performance of resnet18 model")
ax_6_2.plot(range(num_epochs), train_accuracy_list_8, label = 'train_acc')
ax_6_2.legend()
ax_6_2.plot(range(num_epochs), valid_accuracy_list_8, label = 'valid_acc')
ax_6_2.legend()
ax_6_2.set_xlabel("Epoch")
ax_6_2.set_ylabel("accuracy")
# ax_6_2.set_title("Performance of resnet18 model")
plt.title("Performance of resnet18 model with a classification layer composed of two layers and hidden dim = 256")
plt.show()

train_loss_list_9 = []
valid_loss_list_9 = []
train_accuracy_list_9 = []
valid_accuracy_list_9 = []
test_loss_list_9 = []
test_accuracy_list_9 = []
hidden_dim = 1024
num_epochs = 10

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    epoch_time = []  # we'll keep track of the time needed for each epoch
    
    

    for epoch in range(num_epochs):
        epoch_start = time.time()
        print("Epoch {}/{}".format(epoch + 1, num_epochs))
        print("-" * 10)

        # Each epoch has a training and validation phase
        for phase in ["train", "val"]:
            if phase == "train":
                scheduler.step()
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # Forward
                # Track history if only in training phase
                with torch.set_grad_enabled(phase == "train"):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

                # Statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
         
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            

            if phase == "train" :
                train_loss_list_9.append(epoch_loss)
                train_accuracy_list_9.append(epoch_acc)
            elif phase == "val":
                valid_loss_list_9.append(epoch_loss)
                valid_accuracy_list_9.append(epoch_acc
                )

            print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc))

            # Deep copy the model
            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        # Add the epoch time
        t_epoch = time.time() - epoch_start
        epoch_time.append(t_epoch)
        print()

    time_elapsed = time.time() - since
    print(
        "Training complete in {:.0f}m {:.0f}s".format(
            time_elapsed // 60, time_elapsed % 60
        )
    )
    print("Best val Acc: {:4f}".format(best_acc))

    # Load best model weights
    model.load_state_dict(best_model_wts)
    return model, epoch_time


def test_model (model, criterion):
    since = time.time()

    running_loss = 0.0
    running_corrects = 0
    model.eval()
    for inputs, labels in dataloaders["test"] :
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
    
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
        
    test_loss = running_loss/dataset_sizes["test"]
    test_accuracy = running_corrects.double()/dataset_sizes["test"]
    test_loss_list_9.append(test_loss)
    test_accuracy_list_9.append(test_accuracy)

    print("test loss is %2f and test accuracy  is %2f" %(test_loss, test_accuracy ))

    return test_loss, test_accuracy
    


# Download a pre-trained ResNet18 model and freeze its weights
model_N3 = torchvision.models.resnet18(pretrained=True)
for param in model_N3.parameters():
    param.requires_grad = False

# Replace the final fully connected layer
# Parameters of newly constructed modules have requires_grad=True by default

num_ftrs = model_N3.fc.in_features
model_N3.fc = nn.Sequential(
nn.Dropout(0.2),     
nn.Linear(num_ftrs, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 2))

# Send the model to the GPU
model_N3 = model_N3.to(device)
# Set the loss function
criterion = nn.CrossEntropyLoss()

# Observe that only the parameters of the final layer are being optimized
optimizer_conv = optim.SGD(model_N3.fc.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
model_N3, epoch_time = train_model(
    model_N3, criterion, optimizer_conv, exp_lr_scheduler, num_epochs )

test_loss, test_accuracy = test_model(model_N3, criterion)

fig_6 = plt.figure() 
fig_6.set_size_inches(10, 4)
plt.subplots_adjust(wspace=0.2, hspace=0.35)
ax_6_1 = plt.subplot(1,2,1) 
ax_6_2 = plt.subplot(1,2,2)
ax_6_1.plot(range(num_epochs), train_loss_list_9, label = 'train_loss')
ax_6_1.legend()
ax_6_1.plot(range(num_epochs), valid_loss_list_9, label = 'valid_loss')
ax_6_1.legend()
ax_6_1.set_xlabel("Epoch")
ax_6_1.set_ylabel("Loss")
# ax_6_1.set_title("Performance of resnet18 model") 
ax_6_2.plot(range(num_epochs), train_accuracy_list_9, label = 'train_acc')
ax_6_2.legend()
ax_6_2.plot(range(num_epochs), valid_accuracy_list_9, label = 'valid_acc')
ax_6_2.legend()
ax_6_2.set_xlabel("Epoch")
ax_6_2.set_ylabel("accuracy")
# ax_6_2.set_title("Performance of resnet18 model")
plt.title("Performance of resnet18 model with a classification layer composed of two layers and hidden dim = 1024")
plt.show()

def print_size_of_model(model, label=""):
    torch.save(model.state_dict(), "temp.p")
    size = os.path.getsize("temp.p")
    print("model: ", label, " \t", "Size (KB):", size / 1e3)
    os.remove("temp.p")
    return size

print_size_of_model(model_N1, "fp32")
quantized_model_N1 = torch.quantization.quantize_dynamic(model_N1, dtype=torch.qint8)
print_size_of_model(quantized_model_N1, "int8")

test_loss_list_10 = []
test_accuracy_list_10 = []

def test_model (model, criterion):
    since = time.time()

    running_loss = 0.0
    running_corrects = 0
    model.eval()
    for inputs, labels in dataloaders["test"] :
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
    
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
        
    test_loss = running_loss/dataset_sizes["test"]
    test_accuracy = running_corrects.double()/dataset_sizes["test"]
    test_loss_list_10.append(test_loss)
    test_accuracy_list_10.append(test_accuracy)

    print("test loss is %2f and test accuracy  is %2f" %(test_loss, test_accuracy ))

    return test_loss, test_accuracy
    
# Send the model to the GPU
quantized_model_N1 = quantized_model_N1.to(device)
# Set the loss function
criterion = nn.CrossEntropyLoss()
test_loss, test_accuracy = test_model(quantized_model_N1, criterion)

print_size_of_model(model_N2, "fp32")
quantized_model_N2 = torch.quantization.quantize_dynamic(model_N2, dtype=torch.qint8)
print_size_of_model(quantized_model_N2, "int8")

test_loss_list_11 = []
test_accuracy_list_11 = []

def test_model (model, criterion):
    since = time.time()

    running_loss = 0.0
    running_corrects = 0
    model.eval()
    for inputs, labels in dataloaders["test"] :
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
    
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
        
    test_loss = running_loss/dataset_sizes["test"]
    test_accuracy = running_corrects.double()/dataset_sizes["test"]
    test_loss_list_11.append(test_loss)
    test_accuracy_list_11.append(test_accuracy)

    print("test loss is %2f and test accuracy  is %2f" %(test_loss, test_accuracy ))

    return test_loss, test_accuracy
    
# Send the model to the GPU
quantized_model_N2 = quantized_model_N2.to(device)
# Set the loss function
criterion = nn.CrossEntropyLoss()
test_loss, test_accuracy = test_model(quantized_model_N2, criterion)

print_size_of_model(model_N3, "fp32")
quantized_model_N3 = torch.quantization.quantize_dynamic(model_N3, dtype=torch.qint8)
print_size_of_model(quantized_model_N3, "int8")

test_loss_list_12 = []
test_accuracy_list_12 = []

def test_model (model, criterion):
    since = time.time()

    running_loss = 0.0
    running_corrects = 0
    model.eval()
    for inputs, labels in dataloaders["test"] :
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
    
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)
        
    test_loss = running_loss/dataset_sizes["test"]
    test_accuracy = running_corrects.double()/dataset_sizes["test"]
    test_loss_list_12.append(test_loss)
    test_accuracy_list_12.append(test_accuracy)

    print("test loss is %2f and test accuracy  is %2f" %(test_loss, test_accuracy ))

    return test_loss, test_accuracy
    
# Send the model to the GPU
quantized_model_N3 = quantized_model_N3.to(device)
# Set the loss function
criterion = nn.CrossEntropyLoss()
test_loss, test_accuracy = test_model(quantized_model_N3, criterion)