import torch
from torch import nn

import torchvision
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import ToTensor

import matplotlib.pyplot as plt

print(torch.__version__)
print(torchvision.__version__)

# Setup training data
train_data = datasets.FashionMNIST(
    root="data", # where to download data to
    train=True, # do we want training dataset
    download=True, # do we want to download data to computer
    transform=ToTensor(), # how do we want to transform data
    target_transform=None # how do we want to transform labels
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
    target_transform=None
)

len(train_data), len(test_data)

# See the first training example
image, label = train_data[0]
image, label

class_names = train_data.classes
class_names

class_to_idx = train_data.class_to_idx
class_to_idx

# Check the shape
print(f"image shape: {image.shape} -> [color_channels, height, width], label: {class_names[label]}")

import matplotlib.pyplot as plt
image, label = train_data[0]
print(f"Image shape: {image.shape}")
plt.imshow(image.squeeze())
plt.title(label)

plt.imshow(image.squeeze(), cmap="gray")
plt.title(class_names[label])
plt.axis(False)

# Plot more images
torch.manual_seed(42)

fig=plt.figure(figsize=(9,9))
rows, cols = 4, 4
for i in range(1, rows*cols+1):
    random_idx = torch.randint(0, len(train_data), size=[1]).item()
    img, label = train_data[random_idx]
    fig.add_subplot(rows, cols, i)
    plt.imshow(img.squeeze(), cmap="gray")
    plt.title(class_names[label])
    plt.axis(False)

train_data, test_data

from torch.utils.data import DataLoader

# Setup the batch size hyperparameter
BATCH_SIZE = 32

# Turn datasets into iterables (batches)
train_dataloader = DataLoader(
    dataset=train_data,
    batch_size=BATCH_SIZE,
    shuffle=True
)

test_dataloader = DataLoader(
    dataset=test_data,
    batch_size=BATCH_SIZE,
    shuffle=False
)

train_dataloader, test_dataloader

# Let's check out what we've created

print(f"DataLoaders: {train_dataloader, test_dataloader}")
print(f"Length of train_dataloader: {len(train_dataloader)} batches of {BATCH_SIZE}")
print(f"Length of test_dataloader: {len(test_dataloader)} of batch size {BATCH_SIZE}")

train_features_batch, train_labels_batch = next(iter(train_dataloader))
train_features_batch.shape, train_labels_batch.shape

# Show a sample
torch.manual_seed(42)
random_idx = torch.randint(0, len(train_features_batch), size=[1]).item()
img, label = train_features_batch[random_idx], train_labels_batch[random_idx]

plt.imshow(img.squeeze(), cmap="gray")
plt.title(class_names[label])
plt.axis(False)
print(f"Image size: {img.shape}")
print(f"Label: {label}, label size: {label.shape}")

# Create a flatten layer
flatten_model = nn.Flatten()

# Get a single sample
x = train_features_batch[0]

# Flatten the sample
output = flatten_model(x) # perform forward pass

# Print out what happened
print(f"Shape before flattening: {x.shape} -> [color_channels, height, width]")
print(f"Shape after flattening: {output.shape} -> [color_channels, height*width]")

from torch import nn

class FashionMNISTModelV0(nn.Module):
    def __init__(
        self,
        input_shape: int,
        hidden_units: int,
        output_shape: int
    ):
        super().__init__()
        self.layer_stack = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=input_shape, out_features=hidden_units),
            nn.Linear(in_features=hidden_units, out_features=output_shape)
        )
        
    def forward(self, x): 
        return self.layer_stack(x)

torch.manual_seed(42)

# Setup model with input parameters
model_0 = FashionMNISTModelV0(
    input_shape=784,
    hidden_units=10,
    output_shape=len(class_names)
).to("cpu")

model_0

dummy_x = torch.rand([1, 1, 28, 28])
model_0(dummy_x)

import requests
from pathlib import Path

# Download helper functions from Learn Pytorch repo
if Path("helper_functions.py").is_file():
    print("already exists")
else:
    print("downloading file")
    request = requests.get("https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/main/helper_functions.py")
    with open("helper_functions.py", "wb") as f:
        f.write(request.content)

# Import accuracy metric
from helper_functions import accuracy_fn

# Setup loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_0.parameters(), lr=0.1)

from timeit import default_timer as timer

def print_train_time(start: float, end: float, device: torch.device = None):
    """
    Prints difference between start and end time
    """
    
    total_time = end - start
    print(f"Train time on {device}: {total_time:.3f} seconds")
    return total_time

start_time = timer()
end_time = timer()
print_train_time(start=start_time, end=end_time, device="cpu")

# Import tqdm for progress bar
from tqdm.auto import tqdm

# Set the seed and start the timer
torch.manual_seed(42)
train_time_start_on_cpu = timer()

# Set the number of epochs (we'll keep it small for faster training loop)
epochs = 3

# Create training and test loop
for epoch in tqdm(range(epochs)):
    print(f"Epoch: {epoch}\n------")
    # Training
    train_loss = 0
    
    # Add a loop to loop through the training batches
    for batch, (X, y) in enumerate(train_dataloader):
        model_0.train()
        # 1. Forward pass
        y_pred = model_0(X)
        
        # 2. Calculate the loss
        loss = loss_fn(y_pred, y)
        train_loss += loss # accumulate train loss
        
        # 3. Optimizer zero grad
        optimizer.zero_grad()
        
        # 4. Loss bachward
        loss.backward()
        
        # 5. Optimizer step
        optimizer.step()
        
        # Print out what's happening
        if batch % 400 == 0:
            print(f"Looked at {batch * len(X)}/{len(train_dataloader.dataset)} examples")
            
    # Devide total train loss by length of train dataloader
    train_loss /= len(train_dataloader)
    
    ### Testing
    test_loss, test_acc = 0, 0
    model_0.eval()
    with torch.inference_mode():
        for X_test, y_test in test_dataloader:
            # 1. froward pass
            test_pred = model_0(X_test)
            
            # 2. Calculate the loss (accumulatevly)
            test_loss += loss_fn(test_pred, y_test)
            
            # 3. Calculate accuracy
            test_acc += accuracy_fn(y_true=y_test, y_pred=test_pred.argmax(dim=1))
            
        # Calculate the test loss average per batch
        test_loss /= len(test_dataloader)
        
        # Calculate the test acc average per batch
        test_acc /= len(test_dataloader)
    
    # Print out what's happening
    print(f"\nTrain loss: {train_loss:.4f} | test loss: {test_loss:.4f} | test accuracy: {test_acc:.2f}")
    
# Calculate train time
train_time_end_on_cpu = timer()
total_train_time_model_0 = print_train_time(
    start=train_time_start_on_cpu,
    end=train_time_end_on_cpu,
    device=str(next(model_0.parameters()).device)
)

torch.manual_seed(42)
def eval_model(model: torch.nn.Module,
              data_loader: torch.utils.data.DataLoader,
              loss_fn: torch.nn.Module,
              accuracy_fn,
              device = "cpu"):
    """
    Returns a dictionary containing the results of model predicting on data_loader
    """
    
    loss, acc = 0, 0
    model.eval()
    with torch.inference_mode():
        for X, y in tqdm(data_loader):
            # Make our data device agnostic
            X, y = X.to(device), y.to(device)
            # Make predictions
            y_pred = model(X)
            
            # Accumulate the loss and acc values per batch
            loss += loss_fn(y_pred, y)
            acc += accuracy_fn(y_true=y,
                              y_pred=y_pred.argmax(dim=1))
            
        #Scale loss and acc to find the average loss/acc per batch
        loss /= len(data_loader)
        acc /= len(data_loader)
    
    return {"model_name": model.__class__.__name__,
            "model_loss": loss.item(),
            "model_acc": acc}

# Calculate model 0 results on test dataset
model_0_results = eval_model(model=model_0,
                            data_loader=test_dataloader,
                            loss_fn=loss_fn,
                            accuracy_fn=accuracy_fn,
                            device="cpu")
model_0_results

device = "cuda" if torch.cuda.is_available() else "cpu"
device

# Create a model with linear and non-linear layers
class FashionMNISTModelV1(nn.Module):
    def __init__(self,
                input_shape: int,
                hidden_units: int,
                output_shape: int):
        super().__init__()
        self.layer_stack = nn.Sequential(
            nn.Flatten(), # Flatten input into a single vector
            nn.Linear(in_features=input_shape, out_features=hidden_units),
            nn.ReLU(),
            nn.Linear(in_features=hidden_units, out_features=hidden_units),
            nn.ReLU(),
            nn.Linear(in_features=hidden_units, out_features=output_shape)
        )
        
    def forward(self, x):
        return self.layer_stack(x)

device

!nvidia-smi

# Create an instance of model
model_1 = FashionMNISTModelV1(input_shape=28*28,
                             hidden_units=30,
                             output_shape=len(class_names)).to(device)

model_1

next(model_1.parameters()).device

next(model_0.parameters()).device

from helper_functions import accuracy_fn
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_1.parameters(), lr=0.1)

def train_step(model: torch.nn.Module,
               data_loader: torch.utils.data.DataLoader,
               loss_fn: torch.nn.Module,
               optimizer: torch.optim.Optimizer,
               accuracy_fn,
               device: torch.device = device):
    """
    Performs a training with model trying to learn a data_loader
    """
    # Training
    train_loss, train_acc = 0, 0
    
    # Put model into training mode    
    model.train()
    
    # Add a loop to loop through the training batches
    for batch, (X, y) in enumerate(data_loader):
        # Put data on target device
        X, y = X.to(device), y.to(device)
        
        # 1. Forward pass
        y_pred = model(X)
        
        # 2. Calculate the loss and accuracy
        loss = loss_fn(y_pred, y)
        train_loss += loss # accumulate train loss
        train_acc += accuracy_fn(y_true=y, y_pred=y_pred.argmax(dim=1))
        
        # 3. Optimizer zero grad
        optimizer.zero_grad()
        
        # 4. Loss bachward
        loss.backward()
        
        # 5. Optimizer step
        optimizer.step()
            
    # Devide total train loss and accuracy by length of train dataloader
    train_loss /= len(data_loader)
    train_acc /= len(data_loader)
    print(f"Train loss: {train_loss:.3f} | Train accuracy: {train_acc:.2f}%")
    

def test_step(model: torch.nn.Module,
               data_loader: torch.utils.data.DataLoader,
               loss_fn: torch.nn.Module,
               accuracy_fn,
               device: torch.device = device):
    """
    Performs a testing with data_loader trying to evaluate the model
    """
    # Training
    test_loss, test_acc = 0, 0
    
    # Put model into testing mode    
    model.eval()
    
    # Turn on inference mode context manager
    with torch.inference_mode():
        # Add a loop to loop through the training batches
        for X, y in data_loader:
            # Put data on target device
            X, y = X.to(device), y.to(device)
        
            # 1. Forward pass
            y_pred = model(X)
        
            # 2. Calculate the loss and accuracy
            loss = loss_fn(y_pred, y)
            test_loss += loss # accumulate train loss
            test_acc += accuracy_fn(y_true=y, y_pred=y_pred.argmax(dim=1))
            
        # Devide total train loss and accuracy by length of test dataloader
        test_loss /= len(data_loader)
        test_acc /= len(data_loader)
        print(f"Test loss: {test_loss:.3f} | Test accuracy: {test_acc:.2f}%\n")

# Import tqdm for progress bar
from tqdm.auto import tqdm

# Set the seed and start the timer
torch.manual_seed(42)
train_time_start_on_gpu = timer()

# Set the number of epochs (we'll keep it small for faster training loop)
epochs = 3

# Create training and test loop
for epoch in tqdm(range(epochs)):
    print(f"Epoch: {epoch}\n------")
    train_step(
        model = model_1,
        data_loader = train_dataloader,
        loss_fn = loss_fn,
        optimizer = optimizer,
        accuracy_fn = accuracy_fn,
        device = device
    )
    test_step(
        model = model_1,
        data_loader = test_dataloader,
        loss_fn = loss_fn,
        accuracy_fn = accuracy_fn,
        device = device
    )
    
# Calculate train time
train_time_end_on_gpu = timer()
total_train_time_model_1 = print_train_time(
    start=train_time_start_on_gpu,
    end=train_time_end_on_gpu,
    device=str(next(model_1.parameters()).device)
)

model_0_results

# Get model_1 results dictionary
model_1_results = eval_model(model=model_1,
                            data_loader=test_dataloader,
                            loss_fn=loss_fn,
                            accuracy_fn=accuracy_fn,
                            device=device)
model_1_results

model_0_results

# Create a convolutional neural network
    
class FashionMNISTModelV2(nn.Module):
    """
    Model architecture that replicates the TinyVGG
    model from CNN explainer website
    """
    
    def __init__(self, input_shape: int,
                hidden_units: int,
                output_shape: int):
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            # Create a conv layer
            nn.Conv2d(in_channels = input_shape,
                      out_channels= hidden_units,
                     kernel_size=3,
                     stride=1,
                     padding=1), # values we can set ourselves in our NN's are called hyperparameters
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units,
                     out_channels=hidden_units,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(in_channels=hidden_units,
                     out_channels=hidden_units,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units,
                     out_channels=hidden_units,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=hidden_units*7*7, # there is a trick to calculate this
                     out_features=output_shape)
        )
        
    def forward(self, x):
        x = self.conv_block_1(x)
        #print(x.shape)
        x = self.conv_block_2(x)
        #print(x.shape)
        x = self.classifier(x)
        #print(x.shape)
        return x

torch.manual_seed(42)
model_2 = FashionMNISTModelV2(input_shape=1,
                             hidden_units=10,
                             output_shape=len(class_names)).to(device)
model_2

plt.imshow(image.squeeze(), cmap="gray")

rand_mage_tensor = torch.randn(size=(1, 28, 28)).to(device)
# Pass image through model
model_2(rand_mage_tensor.unsqueeze(0))

torch.manual_seed(42)

# Create a batch of images
images = torch.randn(size=(32, 3, 64, 64))
test_image = images[0]

print(f"Image batch shape: {images.shape}")
print(f"Single image shape: {test_image.shape}")

# Create a single conv2d layer
conv_layer = nn.Conv2d(in_channels=3,
                      out_channels=10,
                      kernel_size=(3,3),
                      stride=1,
                      padding=1)

# Pass the data through the convolutional layer
conv_output = conv_layer(test_image.unsqueeze(0))
conv_output.shape

test_image.shape

# Print out the original image shape without unsqueeze dimension
print(f"Test image original shape: {test_image.shape}")
print(f"Test image with unsqueezed dimension: {test_image.unsqueeze(0).shape}")

# Create a sample nn.MaxPool2d layer
max_pool_layer = nn.MaxPool2d(kernel_size=2)

# Pass data through just the conv layer
tst_image_through_conv = conv_layer(test_image.unsqueeze(dim=0))
print(f"Shape after going through conv layer: {tst_image_through_conv.shape}")

# Pass data through the max pool layer
test_image_through_conv_and_max_pool = max_pool_layer(tst_image_through_conv)
print(f"Shape after going through conv layer and max pool layer: {test_image_through_conv_and_max_pool.shape}")

torch.manual_seed(42)
# Create a random tensor with a similar number of dimensions to our image
random_tensor = torch.randn(size=(1, 1, 2, 2))
print(f"Original Tensor:\n {random_tensor}")
print(f"Original Tensor shape: {random_tensor.shape}")

# Create max pool layer
max_pool_layer = nn.MaxPool2d(kernel_size=2)

# Pass the random tensor through the max pool layer
random_tensor_through_max_pool = max_pool_layer(random_tensor)
print(f"Max pool tensor:\n {random_tensor_through_max_pool}")
print(f"Max pool tensor shape: {random_tensor_through_max_pool.shape}")

# Setup loss function/eval metrics/optimizer
from helper_functions import accuracy_fn

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_2.parameters(), lr=0.1)

torch.manual_seed(42)
torch.cuda.manual_seed(42)

# Measure time
from timeit import default_timer as timer
train_time_start_model_2 = timer()

# Train and test model
epochs = 3
for epoch in tqdm(range(epochs)):
    print(f"Epoch: {epoch}\n-------")
    train_step(model=model_2,
              data_loader=train_dataloader,
              loss_fn=loss_fn,
              optimizer=optimizer,
              accuracy_fn=accuracy_fn,
              device=device)
    test_step(model=model_2,
             data_loader=test_dataloader,
             loss_fn=loss_fn,
             accuracy_fn=accuracy_fn,
             device=device)
    
train_time_end_model_2 = timer()
total_train_time_model_2 = print_train_time(
    start=train_time_start_model_2,
    end=train_time_end_model_2,
    device=device
)

# Get model_2 results
model_2_results = eval_model(model=model_2,
                            data_loader=test_dataloader,
                            loss_fn=loss_fn,
                            accuracy_fn=accuracy_fn,
                            device=device)
model_2_results

import pandas as pd
compare_results = pd.DataFrame([model_0_results,
                               model_1_results,
                               model_2_results])
compare_results

# Add training time to results comparison
compare_results["training_time"] = [total_train_time_model_0,
                                   total_train_time_model_1,
                                   total_train_time_model_2]
compare_results

# Visualize our model results
compare_results.set_index("model_name")["model_acc"].plot(kind="barh")
plt.xlabel("accuracy %")
plt.ylabel("model")

def make_predictions(model: torch.nn.Module,
                    data: list,
                    device: torch.device = device):
    pred_probs = []
    model.to(device)
    model.eval()
    with torch.inference_mode():
        for sample in data:
            # Prepare the sample (add a batch dimension and pass to target device)
            sample = torch.unsqueeze(sample, dim=1).to(device)
            
            # Forward pass (model output row logits)
            pred_logits = model(sample)
            
            # Get prediction pobability (logit -> prediction probabilities)
            pred_prob = torch.softmax(pred_logits.squeeze(), dim=0)
            
            # Get pred_prob off the GPU for futher calculations
            pred_probs.append(pred_prob.cpu())
        
    # Stack the pred probs to turn list into tensor    
    return torch.stack(pred_probs)   

import random
#random.seed(42)
test_samples = []
test_labels = []
for sample, label in random.sample(list(test_data), k=9):
    test_samples.append(sample)
    test_labels.append(label)
    
# View the first sample shape
test_samples[0].shape

plt.imshow(test_samples[0].squeeze(), cmap="gray")
plt.title(class_names[test_labels[0]])

# Make predictions
pred_probs = make_predictions(model=model_2,
                             data=test_samples)

# View first two prediction probablities
pred_probs[:2]

# Convert prediction probabilities to labels
pred_classes = pred_probs.argmax(dim=1)
pred_classes

test_labels

# Plot predictions
plt.figure(figsize=(9, 9))
nrows = 3
ncols = 3
for i, sample in enumerate(test_samples):
    # Create subplot
    plt.subplot(nrows, ncols, i+1)
    
    # Plot the target image
    plt.imshow(sample.squeeze(), cmap="gray")
    
    # Find the prediction (in text form e.g "Sandal")
    pred_label = class_names[pred_classes[i]]
    
    # Get the truth label (in text form)
    truth_label = class_names[test_labels[i]]
    
    # Create a title for the plot
    title_text = f"Pred: {pred_label} | Truth: {truth_label}"
    
    # Check for equality between pred and truth and change color of title text
    if pred_label == truth_label:
        plt.title(title_text, fontsize=10, c="g")
    else:
        plt.title(title_text, fontsize=10, c="r")
    plt.axis(False)

from tqdm.auto import tqdm

# 1. Make predictions with trained model
y_preds = []
model_2.eval()
with torch.inference_mode():
    for X, y in tqdm(test_dataloader, desc="Making predictions..."):
        # Send the data and targets to target device
        X, y = X.to(device), y.to(device)
        
        # Do the forward pass
        y_logit = model_2(X)
        
        # Turn predictions from logits -> prediction probabilities -> prediction labels
        y_pred = torch.softmax(y_logit.squeeze(), dim=0).argmax(dim=1)
        
        # Put predictions on cpu for evaluations
        y_preds.append(y_pred.cpu())

# Concatenate list of predictions into tensor
print(y_preds[0].shape)
y_pred_tensor = torch.cat(y_preds)
print(y_pred_tensor.shape)
y_pred_tensor[:10]

import torchmetrics, mlxtend
print(f"mlxtend version: {mlxtend.__version__}")
assert int(mlxtend.__version__.split(".")[1]) >= 19, "mlxtend version should be 0.19.0 or higher"

from torchmetrics import ConfusionMatrix
from mlxtend.plotting import plot_confusion_matrix

# 2. Setup confusion instance and compare predictions to targets
confmat = ConfusionMatrix(task='multiclass', num_classes=len(class_names))
confmat_tensor = confmat(preds=y_pred_tensor,
                        target=test_data.targets)

# 3. Plot the confusion matrix
fig, ax = plot_confusion_matrix(
    conf_mat=confmat_tensor.numpy(), # matplotlib likes working with numpy
    class_names=class_names,
    figsize=(10, 7)
)

confmat_tensor

from pathlib import Path

# Create model directory path
MODEL_PATH = Path("models")
MODEL_PATH.mkdir(parents=True,
                exist_ok=True)

# Create model save
MODEL_NAME = "03_pytorch_computer_vision_model_2.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME
MODEL_SAVE_PATH

# Save the model state dict
torch.save(obj=model_2.state_dict(), f=MODEL_SAVE_PATH)

# Create new instance 
torch.manual_seed(42)

loaded_model_2 = FashionMNISTModelV2(input_shape=1,
                                    hidden_units=10,
                                    output_shape=len(class_names))

# Load in the save state_dict()
loaded_model_2.load_state_dict(torch.load(f=MODEL_SAVE_PATH))
loaded_model_2 = loaded_model_2.to(device)

model_2_results

# Evaluate loaded model
torch.manual_seed(42)

loaded_model_2_results = eval_model(
    model=loaded_model_2,
    data_loader=test_dataloader,
    loss_fn=loss_fn,
    accuracy_fn=accuracy_fn,
    device=device
)
loaded_model_2_results

# Check if model results are close to each other
torch.isclose(torch.tensor(model_2_results["model_loss"]),
             torch.tensor(loaded_model_2_results["model_loss"]),
             atol=1e-02) # tollerance

import torchvision
from torchvision import datasets
from torchvision.transforms import ToTensor

mnist_train = datasets.MNIST(
    root="data_mnist",
    train=True,
    download=True,
    transform=ToTensor(),
    target_transform=None
)

mnist_test = datasets.MNIST(
    root="data_mnist",
    train=False,
    download=True,
    transform=ToTensor(),
    target_transform=None
)

len(mnist_train), len(mnist_test)

import matplotlib.pyplot as plt
import torch

n_row = 1
n_col = 5
_, axs = plt.subplots(n_row, n_col, figsize=(12, 12))
axs = axs.flatten()
for i, ax in zip(range(5), axs):
    random_train_number = torch.randint(low=0, high=len(mnist_train), size=()).item()
    image, label = mnist_train[random_train_number]
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(f"Number {label}")
plt.show()

from torch.utils.data import DataLoader

train_dataloader = DataLoader(mnist_train,
                            batch_size=32,
                            shuffle=True)

test_dataloader = DataLoader(mnist_test,
                            batch_size=32,
                            shuffle=False)

len(train_dataloader), len(test_dataloader), mnist_train.classes

from torch import nn

class MNISTModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer_1 = nn.Sequential(
            nn.Conv2d(in_channels=1,
                     out_channels=10,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=10,
                     out_channels=10,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.layer_2 = nn.Sequential(
            nn.Conv2d(in_channels=10,
                     out_channels=10,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=10,
                     out_channels=10,
                     kernel_size=3,
                     stride=1,
                     padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=10*7*7,
                     out_features=len(mnist_train.classes))
        )
        
    def forward(self, x):
        x = self.layer_1(x)
        x = self.layer_2(x)
        x = self.classifier(x)
        return x
    
model_e = MNISTModel().to("cuda")
model_e

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_e.parameters(), lr=0.1)

from tqdm.auto import tqdm
from timeit import default_timer as timer

def train_and_test(model, device, loss_fn, optimizer):
    start_time = timer()
    torch.manual_seed(42)
    torch.cuda.manual_seed(42)
    epochs = 3
    
    for epoch in tqdm(range(epochs), "Training model"):        
        loss_epoch, accuracy_epoch = 0, 0
        test_loss_epoch, test_accuracy_epoch = 0, 0
        
        model.train()
        for batch, (X, y) in enumerate(train_dataloader):
            X, y = X.to(device), y.to(device)
            y_logits = model(X)
            
            loss = loss_fn(y_logits, y)
            accuracy = accuracy_fn(y_pred=y_logits.argmax(dim=1),
                                  y_true=y)
            loss_epoch += loss
            accuracy_epoch += accuracy
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        loss_epoch /= len(train_dataloader)
        accuracy_epoch /= len(train_dataloader)
        
        model.eval()
        with torch.inference_mode():
            for batch, (X, y) in enumerate(test_dataloader):
                X, y = X.to(device), y.to(device)
                y_logits = model(X)
                
                loss_test = loss_fn(y_logits, y)
                accuracy_test = accuracy_fn(y_pred=y_logits.argmax(dim=1), y_true=y)
                test_loss_epoch += loss_test
                test_accuracy_epoch += accuracy_test
                
            test_loss_epoch /= len(test_dataloader)
            test_accuracy_epoch /= len(test_dataloader)
        
        print(f"Loss: {loss_epoch:.2f}. Accuracy: {accuracy_epoch:.2f}. Test Loss: {test_loss_epoch:.2f}. Test Accuracy: {test_accuracy_epoch:.2f}")    
        
    end_time = timer()
    print(f"Time spend on training on {device} is {end_time - start_time} seconds")
            

train_and_test(model_e, "cuda", loss_fn, optimizer)

model_e_cpu = MNISTModel().to("cpu")
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_e_cpu.parameters(), lr=0.1)
train_and_test(model_e_cpu, "cpu", loss_fn, optimizer)

n_row = 1
n_col = 5
model_e = model_e.to("cpu")
_, axs = plt.subplots(n_row, n_col, figsize=(12, 12))
axs = axs.flatten()
for i, ax in zip(range(5), axs):
    random_train_number = torch.randint(low=0, high=len(mnist_test), size=()).item()
    image, label = mnist_train[random_train_number]
    y_logit = model_e(image.unsqueeze(dim=0))
    ax.imshow(image.squeeze(), cmap="gray")
    ax.set_title(f"{label}/{y_logit.argmax(dim=1).item()}")
plt.show()

y_preds_numbers = []
model_e = model_e.to(device)
model_e.eval()
with torch.inference_mode():
    for X, y in tqdm(test_dataloader, desc="Making predictions..."):
        # Send the data and targets to target device
        X, y = X.to(device), y.to(device)
        
        # Do the forward pass
        y_logit = model_e(X)
        
        # Turn predictions from logits -> prediction probabilities -> prediction labels
        y_pred = torch.softmax(y_logit.squeeze(), dim=0).argmax(dim=1)
        
        # Put predictions on cpu for evaluations
        y_preds_numbers.append(y_pred.cpu())

# Concatenate list of predictions into tensor
print(y_preds_numbers[0].shape)
y_pred_tensor_numbers = torch.cat(y_preds_numbers)
print(y_pred_tensor_numbers.shape)
y_pred_tensor_numbers[:10]

from torchmetrics import ConfusionMatrix
from mlxtend.plotting import plot_confusion_matrix

# 2. Setup confusion instance and compare predictions to targets
confmat = ConfusionMatrix(task='multiclass', num_classes=len(mnist_test.classes))
confmat_tensor = confmat(preds=y_pred_tensor_numbers,
                        target=mnist_test.targets)

# 3. Plot the confusion matrix
fig, ax = plot_confusion_matrix(
    conf_mat=confmat_tensor.numpy(), # matplotlib likes working with numpy
    class_names=mnist_test.classes,
    figsize=(10, 7)
)

with torch.inference_mode():
    random_tensor = torch.rand(size=(1, 3, 64, 64))
    conv_layer = nn.Conv2d(in_channels=3,
                      out_channels=2,
                      kernel_size=64,
                      padding=0,
                      stride=1)

result = conv_layer(random_tensor)
result.shape, result