# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="Wb5cCNGSucOV"
# ## Imports

# + id="7Bq9xYnuubGi"
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import numpy as np
import torch.nn.functional as F
from torchsummary import summary

# + [markdown] id="9ssymKcCujFZ"
# ## Loading the MNIST Dataset

# + id="tryfs0quulRt"
train_dataset = torchvision.datasets.MNIST(
                    root='.',
                    train=True,
                    transform=transforms.ToTensor(), 
                    download=True)

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

# + [markdown] id="xyFEJpkduuGS"
#  ## Generate random numbers

# + id="i213RXOSupry"
# Train Set, 60000 Samples
number_train = torch.randint(0, 9, (60000,)) # Random Digit i.e. Second Input
sum_train = number_train + train_dataset.targets # Sum of Random Digit and MNIST i.e. Second Output
number_train = torch.nn.functional.one_hot(number_train, num_classes=10)

# Test Set, 10000 Samples
number_test = torch.randint(0, 9, (10000,)) # Random Digit i.e. Second Input
sum_test = number_test + test_dataset.targets # Sum of Random Digit and MNIST i.e. Second Output
number_test = torch.nn.functional.one_hot(number_test, num_classes=10)

# + [markdown] id="C_69MdhcvFOS"
# ## Re-creating Dataset

# + id="F6Ui8J4MvG39"
# Flattens the MNIST images
train_x = train_dataset.data.reshape(60000, 784).float()
test_x = test_dataset.data.reshape(10000, 784).float()

# Creats the Dataset with two inputs and two outputs
train_ds = torch.utils.data.TensorDataset(train_x, number_train, train_dataset.targets, sum_train)
test_ds = torch.utils.data.TensorDataset(test_x, number_test, test_dataset.targets, sum_test)

# + [markdown] id="G-F0OAlBvIfb"
# ## DataLoader

# + id="SpqFLHofvL-A"
batch_size = 32
train_loader = torch.utils.data.DataLoader(
                        dataset=train_ds,
                        batch_size=batch_size,
                        shuffle=True
                        )

test_loader = torch.utils.data.DataLoader(
                        dataset=test_ds,
                        batch_size=batch_size,
                        shuffle=False # Not necessary!
                        )


# + [markdown] id="ZBBXOpDsvOVm"
# ## Model

# + id="TgOtM7hEvPhu"
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.lineara1 = nn.Linear(784, 30) # Flattened MNIST as Input: 28 * 28
        self.relu = nn.ReLU()
        self.selu = nn.SELU()
        self.lineara2 = nn.Linear(30, 30)
        self.lineara3 = nn.Linear(30,10) # 10 Classes for First Output: 0-9

        self.linearb1 = nn.Linear(40, 60) # 10 for random digit, 30 from output of lineara1 layer
        self.linearb2 = nn.Linear(60, 30)
        self.linearb3 = nn.Linear(30,18) # 18 Classes for Second Output: 0-18
            
    def forward(self, Xa, Xb):

        out = self.lineara1(Xa)
        out1 = self.selu(out)
        out = self.lineara2(out1)
        out = self.selu(out)
        out = self.lineara3(out)
        outa = self.selu(out)


        inb = torch.cat((out1,Xb), dim=-1) # Gets input from the first layer after the input layer
        outb = self.linearb1(inb)
        outb = self.relu(outb)
        outb = self.linearb2(outb)
        outb = self.relu(outb)

        outb = self.linearb3(outb)
        outb = self.relu(outb)
        return outa, outb

# Instantiate the Model
model = Model()

# + [markdown] id="CXu7e9qSvUQx"
# ## Move Model to GPU (Required Condition by Assignment!)

# + colab={"base_uri": "https://localhost:8080/"} id="fYPA-pEIvaUS" outputId="3c5e8090-6b22-477f-9c55-ac438cb60a22"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
model.to(device)

# + [markdown] id="lRI8imRdvcHv"
# ## Model Summary

# + colab={"base_uri": "https://localhost:8080/"} id="qa4iXVqLvdpW" outputId="2847d41a-529d-418c-832c-d61b50818ac9"
summary(model, [(1,784), (1,10)])

# + [markdown] id="kO9dQJJfvfPc"
# ## Loss and Optimizer

# + id="o79IzzV7vhae"
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# + [markdown] id="B6ClCvKqwO4X"
# ## Training Loop

# + colab={"base_uri": "https://localhost:8080/"} id="gpql5R5EwQEd" outputId="dc749db5-e7e8-4f02-960a-2f538677a744"

n_epochs = 20



for epoch in range(n_epochs):
    train_loss = []
    n_correcta = 0.
    n_correctb = 0.  
    n_total = 0.

    for inputs, numbers, targets, result in train_loader:

        # Move data to GPU
        inputs, numbers, targets, result = inputs.to(device), numbers.to(device), targets.to(device), result.to(device)
        
        # zero the gradient
        optimizer.zero_grad()
        
        # forward pass
        outputa, outputb = model(inputs, numbers)
        lossa = loss_fn(outputa, targets)
        lossb = loss_fn(outputb, result)

        loss = lossa + lossb # Adds loss values from both outputs

        # get predictiona
        _, predictiona = torch.max(outputa, 1)
        _, predictionb = torch.max(outputb, 1)

        # update counts
        n_correcta += (predictiona == targets).sum().item()
        n_correctb += (predictionb == result).sum().item()
        n_total += targets.shape[0]
 
        # backward pass and optimize
        loss.backward()
        optimizer.step()
        
        train_loss.append(loss.item())
        
    train_loss = np.mean(train_loss)
    train_acc_a = n_correcta / n_total * 100
    train_acc_b = n_correctb / n_total * 100

    
    test_loss = []
    n_correcta = 0.
    n_correctb = 0. 
    n_total = 0.
    
    for inputs, numbers, targets, result in test_loader:

        # Move data to GPU
        inputs, numbers, targets, result = inputs.to(device), numbers.to(device), targets.to(device), result.to(device)
        
        # forward pass
        outputa, outputb = model(inputs, numbers)
        lossa = loss_fn(outputa, targets)
        lossb = loss_fn(outputb, result)

        loss = lossa + lossb # Adds loss values from both outputs
                
         # get predictions
        _, predictiona = torch.max(outputa, 1) 
        _, predictionb = torch.max(outputb, 1)
        
        
        # update counts
        n_correcta += (predictiona == targets).sum().item()
        n_correctb += (predictionb == result).sum().item()

        n_total += targets.shape[0]

        test_loss.append(loss.item())
        
    test_loss = np.mean(test_loss)
    test_acc_a = n_correcta / n_total * 100
    test_acc_b = n_correctb / n_total * 100
    

    print(f'Epoch: {epoch+1}/{n_epochs}, Train Accuracy [MNIST]: {train_acc_a:.2f}%,  Train Accuracy [Sum]: {train_acc_b:.2f}%, Train Loss: {train_loss:.4f}, Test Accuracy [MNIST]: {test_acc_a:.2f}%,  Test Accuracy [Sum]: {test_acc_b:.2f}%, Train Loss: {test_loss:.4f}')

StarcoderdataJupyterScriptsDedupFiltered