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

# + [markdown] id="bWEEBnVC-Irv"
# # Home Assignment No. 2 - part two
#
# To solve this task, you will write a lot of code to try several machine learning methods for classification and regression.
# * You are **HIGHLY RECOMMENDED** to read relevant documentation, e.g. for [python](https://docs.python.org/3/), [numpy](https://docs.scipy.org/doc/numpy/reference/), [matlpotlib](https://matplotlib.org/) and [sklearn](https://scikit-learn.org/stable/). Also remember that seminars, lecture slides, [Google](http://google.com) and [StackOverflow](https://stackoverflow.com/) are your close friends during this course (and, probably, whole life?).
#
# * If you want an easy life, you have to use **BUILT-IN METHODS** of `sklearn` library instead of writing tons of your own code. There exists a class/method for almost everything you can imagine (related to this homework).
#
# * You have to write **CODE** directly inside specified places marked by comments: **BEGIN/END Solution**. Do not create new cells.
#
# * In some problems you are asked to provide a short discussion of the results. For that find the specific place marked via **Your text answer: \<write your answer\>**.
#
# * For every separate problem or subproblem (if specified) you can get only 0 points or maximal points for this problem. There are **NO INTERMEDIATE scores**. So make sure that you did everything required in the task.
#
# * Your **SOLUTION** notebook **MUST BE REPRODUCIBLE**, i.e., if the reviewer decides to restart the notebook and run all cells, after all the computation he will obtain exactly the same solution (with all the corresponding plots) as in your uploaded notebook. For this purpose, we suggest fixing random `seed` or (better) define `random_state=` inside every algorithm that uses some pseudorandomness.
#
# * Your code must be clear to the reviewer. For this purpose, try to include necessary comments inside the code. But remember: **GOOD CODE MUST BE SELF-EXPLANATORY** without any additional comments.
#
# * Many `sklearn` algorithms support multithreading (Ensemble Methods, Cross-Validation, etc.). Check if the particular algorithm has `n_jobs` parameters and set it to `-1` to use all the cores.
#

# + [markdown] id="ddR3sf3P82Ht"
# ## Task 6. Deep ANNs. (3 points)
#
# -  **(1 pt.)** Activation functions; **(sub tasks 6.1)**
# -  **(2 pt.)** MNIST classification.  **(sub tasks 6.2)**
#
#
#
# ### Task 6.1  Activation functions. 
# Plot the following [activation functions](https://pytorch.org/docs/master/nn.html#non-linear-activation-functions) using their PyTorch implementation and their derivatives using [autograd](https://pytorch.org/docs/stable/autograd.html) functionality `grad()`:
#    
# 1. Plot `ReLU`, `ELU` ($\alpha = 1$), `Softplus` ($\beta = 1$) and `Sign`, `Sigmoid`, `Softsign`, `Tanh`.
#

# + colab={"base_uri": "https://localhost:8080/", "height": 369} id="hcF-2GHz8wMz" outputId="150fa83f-feb5-4fbd-9aa0-b9d02f47707f"
# %matplotlib inline
import torch.nn.functional as F
import matplotlib.pyplot as plt
import torch

x = torch.arange(-2, 2, .01, requires_grad=True)
x_np = x.detach().numpy()
x.sum().backward() # to create x.grad

f, axes = plt.subplots(2, 2, sharex=True, figsize=(14, 5))
axes[0, 0].set_title('Values')
axes[0, 1].set_title('Derivatives')

for i, function_set in (0, (('ReLU', F.relu), ('ELU', F.elu), ('Softplus', F.softplus))), \
                       (1, (('Sign', torch.sign), ('Sigmoid', torch.sigmoid), ('Softsign', F.softsign), ('Tanh', torch.tanh))):
    for function_name, activation in function_set:
        ### BEGIN Solution

            axes[i, 0].plot(x_np, activation(x).detach().numpy(), label=function_name)
            x.grad.zero_()
            activation(x).sum().backward()
            axes[i, 1].plot(x_np, x.grad.detach().numpy(), label=function_name)
            

        ### END Solution

    axes[i, 0].legend()
    axes[i, 1].legend()

plt.tight_layout()
plt.show()

# + [markdown] id="_misNcjO8wXF"
# Which of these functions may be, and which - definitely, are a poor choise as an activation function in a neural network? Why? Do not forget that output of the current layer servers as an input for the following one. Imagine a situation where we have many layers, what happens with the activation values?
#
#
#
#
#
#

# + id="ribXsHDSmlYN"
# BEGIN SOLUTION (do not delete this comment!)


# *   ReLU            good choice, but have zero grad in big range
# *   ELU             good choice
# *   Softplus        good choice
# *   Sign            bad choice: almost everywhere zero derivative
# *   Sigmoid         bad choice: saturates fast - derivative nonzero in small range
# *   SoftSign        maybe bad choice: saturates but slowly then Sigmoid, Tanh
# *   Tanh            bad choice: saturates fast - derivative nonzero in small range

#END SOLUTION (do not delete this comment!)

# + [markdown] id="sW9OYyIw8wz4"
# ### Task 6.2 MNIST classification. 
#
# At one of the seminars we have discussed an MLP (Multilayer perceptron) with one hidden layer, logistic activation functions and softmax. In this task, you are to:
#
# 1. Implement the MLP modules, including the Softmax cross entropy between `logits` and `labels` using numpy.
#
# 2.  Train your numpy realization of MLP to classify the MNIST from `sklearn.datasets()`. The required accuracy on validation is `> 90%`.
#
# 3.  Compare the acccuracy of classification to your scores from `Part 1` with and without dimensionality reduction. Is this comparison fair?:) Derive the confusion matrix for all digits classes. Which digits are predicted better or worse than others, why?

# + id="RKxe88YT_p9P"
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, accuracy_score

# + id="E7_2lHue_r0x"
# fetch the dataset.
digits, targets = load_digits(return_X_y=True)
digits = digits.astype(np.float32) / 255

digits_train, digits_test, targets_train, targets_test = train_test_split(digits, targets, random_state=0)

train_size = digits_train.shape[0]
test_size = digits_test.shape[0]

input_size = 8*8
classes_n = 10

# + colab={"base_uri": "https://localhost:8080/", "height": 111} id="feWQtoKtn1vV" outputId="47edde13-7736-4ef6-f45a-8a315638e00e"
N = 10

sample_idx = np.random.choice(1797, N, replace=False)
digits_sample = digits[sample_idx]
targets_sample = targets[sample_idx]
f, ax  = plt.subplots(1,10, figsize=(10, 5))

for i in range(N):
    ax[i].imshow(digits_sample[i].reshape(8,8))
    ax[i].set_title('label: '+str(targets_sample[i]))


# + [markdown] id="Pj6EctS6yTJK"
# A short recap on what we are going to achieve here.
# <br>
# 1. Forward pass:
# $$
# h_1 = X\theta_1+\beta_1
# $$
#
# $$
# O_1 = sig(h_1)
# $$
#
# $$
# h_2 = O_1\theta_2+\beta_2
# $$
# $$
# O_2 = softmax(h_2)
# $$
# $$
# Loss = CrossEntropy(O_2, true \space labels)
# $$
#
# 2. Compute gradients: 
#
# To update weights first we need to compute loss gradients with respect to $\theta_1$ and $\theta_2$ and then update both $\theta$ and $\beta$.
#
# $$
# \frac{ \partial{loss} }{\partial{\theta_2}} = \frac{ \partial{loss} }{\partial{O_2}}\frac{ \partial{O_2} }{\partial{h_2}}\frac{ \partial{h_2} }{\partial{\theta_2}}
# $$
# Note, that $\frac{ \partial{h_2} }{\partial{\theta_2}}=O_1$, so we can cache this value during forward pass to speed up our computation.
# $$
# \frac{ \partial{loss} }{\partial{\theta_1}} = \frac{ \partial{loss} }{\partial{O_2}}\frac{ \partial{O_2} }{\partial{h_2}}\frac{ \partial{h_2} }{\partial{O_1}}\frac{ \partial{O_1} }{\partial{h_1}}\frac{ \partial{h_1} }{\partial{\theta_1}}
# $$
# Note, that $\frac{ \partial{h_1} }{\partial{\theta_1}}=X$.
#
# Since we are using sigmoid avtivation function here and
# $$
# \frac{ \partial{sig} }{\partial{h}} = sig(h)(1 - sig(h))
# $$
# It also makes sense to cache sig(h) during forward pass.
#
# 3. Update weights:
#
# $\theta:= \theta -  \frac{ \partial{loss} }{\partial{\theta}}\alpha$, where $\alpha$ is some learning rate.
#
# Note, it was not shown here how to update and compute $\beta$ but you can do it!

# + [markdown] id="CaBenjDI_x6k"
# ### Implement the MLP with backprop

# + id="ffpXAKqQ_vfg"
### YOUR TASK STARTS HERE ###

#Here you should implement by yourself MLP class and its constituents including forward and backward propagation methods

class Linear:
    def __init__(self, input_size, output_size):
        # Trainable parameters of the layer and their gradients
        self.thetas = np.random.randn(input_size, output_size) # the weight matrix of the layer (W)
        self.thetas_grads = np.empty_like(self.thetas) # gradient w.r.t. the weight matrix of the layer
        self.bias = np.random.randn(output_size) # bias terms of the layer (b)
        self.bias_grads = np.empty_like(self.bias) # gradient w.r.t. bias terms of the linear layer

    def forward(self, x): 
        # keep x for backward computation
        self.x = x
        output = np.matmul(x, self.thetas) + self.bias
        return output
    
    def backward(self, output_grad, learning_rate):
        """
        Calculate and return gradient of the loss w.r.t. the input of linear layer given the input x and the gradient 
        w.r.t output of linear layer. You should also calculate and update gradients of layer parameters.
        :param x: np.array, input tensor for linear layer;
        :param output_grad: np.array, grad tensor w.r.t output of linear layer;
        :return: np.array, grad w.r.t input of linear layer
        """
        
        # BEGIN SOLUTION (do not delete this comment!)

        input_grad = output_grad @ self.thetas.T
        # calculate mean of gradients across batch w.r.t weights, bias
        n = output_grad.shape[0]
        self.thetas_grads = self.x.T @ output_grad / n
        self.bias_grads = output_grad.mean(axis=0)
        self.step(learning_rate)

        # END Solution (do not delete this comment!)
        
        return input_grad

    def step(self, learning_rate):
        self.thetas -= self.thetas_grads * learning_rate
        self.bias -= self.bias_grads * learning_rate

class LogisticActivation:
    def __init__(self):
        # the layer has no parameters
        pass 
        
    def sig(self, x):
        return 1/(1 + np.exp(-x))

    def forward(self, x):
        # keep o for backward computation
        self.o = self.sig(x)
        return self.o


    def backward(self,  output_grad, learning_rate):
        """
        Calculate and return the gradient of the loss w.r.t. the input
        of logistic non-linearity (given input x and the gradient 
        w.r.t output of logistic non-linearity).
        
        :param x: np.array, input tensor for logistic non-linearity;
        :param output_grad: np.array, grad tensor w.r.t output of logistic non-linearity;
        :return: np.array, grad w.r.t input of logistic non-linearity
        
        """
        
       # BEGIN SOLUTION (do not delete this comment!)

        o = self.o
        input_grad = o * (1 - o) * output_grad

        ### END Solution (do not delete this comment!)
        
        return input_grad


class MLP:
    def __init__(self, input_size, hidden_layer_size, output_size):
        self.linear1 = Linear(input_size, hidden_layer_size)
        self.activation1 = LogisticActivation()
        self.linear2 = Linear(hidden_layer_size, output_size)
        
    def forward(self, x):
        h1 = self.linear1.forward(x)
        h1a = self.activation1.forward(h1)
        out = self.linear2.forward(h1a)
        return out

    def backward(self, output_grad, learning_rate):
        """
        Calculate and return the gradient of the loss w.r.t. the input of MLP given the input and the gradient 
        w.r.t output of MLP. You should also update gradients of paramerters of MLP layers.
        Hint - you should chain backward operations of modules you have already implemented. You may also
        need to calculate intermediate forward results.
        
        :param x: np.array, input tensor for MLP;
        :param output_grad: np.array, grad tensor w.r.t output of MLP;
        :return: np.array, grad w.r.t input of MLP
        """
 
        # BEGIN SOLUTION (do not delete this comment!)

        linear2_input_grad = self.linear2.backward(output_grad, learning_rate)
        activation1_input_grad = self.activation1.backward(linear2_input_grad, learning_rate)
        out = self.linear1.backward(activation1_input_grad, learning_rate)

        # END Solution (do not delete this comment!)
        

        return  out


# + id="07DUqp86_0To"
# BEGIN SOLUTION (do not delete this comment!)

def softmax_crossentropy_with_logits(logits, reference_answers): 

    reference_answers_ = np.zeros_like(logits)
    I = np.arange(logits.shape[0])
    reference_answers_[I, reference_answers] = 1
    loss = np.sum(reference_answers_ * (-logits + np.log(np.sum(np.exp(logits)))))
    
    ### END Solution
    return loss

def grad_softmax_crossentropy_with_logits(logits, reference_answers):
    
    reference_answers_ = np.zeros_like(logits)
    I = np.arange(logits.shape[0])
    reference_answers_[I, reference_answers] = 1
    grad = logits - reference_answers_

    return grad

# BEGIN Solution


# + colab={"base_uri": "https://localhost:8080/"} id="DWkD2V1y_4QU" outputId="d544a97a-41f5-4e64-afae-00b2737e4167"
np.random.seed(42)

mlp = MLP(input_size=input_size, hidden_layer_size=100, output_size=classes_n)

epochs_n = 100
learning_curve = [0] * epochs_n
test_curve = [0] * epochs_n

x_train = digits_train
x_test = digits_test
y_train = targets_train
y_test = targets_test

learning_rate = 1e-2

for epoch in range(epochs_n):
    
    y_pred = []

    for sample_i in range(train_size):
        x = x_train[sample_i].reshape((1, -1))
        target = np.array([y_train[sample_i]])

        ### BEGIN Solution
        
        # ... perform forward pass and compute the loss
        # ... compute the gradients w.r.t. the input of softmax layer
        # ... perform backward pass
        # ... and update the weights with weight -= grad * learning_rate
        

        logits = mlp.forward(x)
        loss = softmax_crossentropy_with_logits(logits, target)
        logits_grad = grad_softmax_crossentropy_with_logits(logits, target)
        mlp.backward(logits_grad, learning_rate)

        ### END Solution

        y_pred.extend(logits.argmax(1))

    if epoch % 10 == 0:
        
        y_pred_test = []

        for sample_i in range(test_size):
            x = x_test[sample_i].reshape((1, -1))
            target = np.array([y_test[sample_i]])
            
            logits = mlp.forward(x)
            y_pred_test.extend(logits.argmax(1))

        print('Starting epoch {}'.format(epoch), \
              ', Loss : {:.3}'.format(loss), \
              ', Accuracy on train: {:.3}'.format(accuracy_score(y_train, y_pred)), \
              ', Accuracy on test: {:.3}'.format(accuracy_score(y_test, y_pred_test)) )
        


# + colab={"base_uri": "https://localhost:8080/"} id="0DNQhxXaARCy" outputId="df7b0f28-cb53-4e7b-d1a3-dfcd222636eb"
# BEGIN SOLUTION (do not delete this comment!)
     
# confusion matrix
from sklearn.metrics import confusion_matrix

y_pred = np.empty_like(y_test)
for sample_i in range(test_size):
    x = x_test[sample_i].reshape((1, -1))
    target = np.array([y_test[sample_i]])

    logits = mlp.forward(x)
    y_pred[sample_i] = logits.argmax(1)

        
confusion_matrix(y_test.astype(np.int), y_pred.astype(np.int))

# END Solution (do not delete this comment!)

# + [markdown] id="MkSdyrpn8xdE"
# ## Task 7. Autoencoders on tabular data (2 points)
# **From now on we will be using pytorch for all the tasks.**
#
# We will build a latent representation for tabular data with simple Autoencoder (AE). We are going to work with the cancer dataset from scikit-learn package. You are to follow the instructions.  
#
# 1. **(1 pt.)**  Implement AE modules for tabular data. Train AE to get latent representation of the cancer dataset from `sklearn.datasets()`. Use `MSE` loss and get < $0.3$ on validation, with AE "bottleneck" = $2$;   **(sub tasks 7.1 - 7.5)**
#
# 2. **(1 pt.)** Plot the latent representation of whole dataset in 2D, use colors to show object of different classes. **(sub tasks: 7.6)**
#

# + id="Sg5fX833AX9q"
# imports

import torch
import torch.nn as nn
import torch.utils.data as torch_data

import sklearn.datasets as sk_data
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler


# + [markdown] id="AYtA62xA8xgB"
# #### 7.1 Fetch the data. Scale it and split on train and test.

# + colab={"base_uri": "https://localhost:8080/"} id="BinFOZc7Abpx" outputId="fa0f0e2d-5114-4343-916a-39bff7ba0722"

cancer_dset = sk_data.load_breast_cancer()

X_train, X_val, y_train, y_val = train_test_split(cancer_dset['data'], cancer_dset['target'], test_size=0.2, random_state=42)
print('\nTrain size: ', len(X_train))
print('Validation size: ', len(X_val))

scaler = StandardScaler() 
scaler.fit(X_train)

X_train = scaler.transform(X_train)
X_val = scaler.transform(X_val)

print('Features: ', list(cancer_dset['feature_names']))
print('\nShape:', X_train.shape)



# + [markdown] id="x7Dzo8VIAaaf"
# #### 7.2 Let us firstly create the dataset, which we'll be able to use with pytorch dataloader. 
# Implement `__len__` and `__getitem__` methods.

# + id="Vi4Cq7DtAl8u"
### BEGIN Solution

class CancerData(torch_data.Dataset):
    def __init__(self, X, y):
        super(CancerData, self).__init__()
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.float32)
    
    def __len__(self):

        return len(self.X)
    
    def __getitem__(self, idx):
        
        
        return self.X[idx], self.y[idx]


### END Solution

# + colab={"base_uri": "https://localhost:8080/"} id="f4emJDB1ApHh" outputId="d695a094-f426-4cf4-880b-c05037673a1e"
train_dset = CancerData(X_train, y_train) 
val_dset = CancerData(X_val, y_val) 

print(train_dset[5])


# + [markdown] id="ksiBurhhAapc"
# #### 7.3 Now, we'll make a base class for our autoencoder. 
# AE takes as input encoder and decoder (it will be two neural networks). Your task is to implement the forward pass.

# + id="wVlgW3_rAqgu"
class MyFirstAE(nn.Module):
    def __init__(self, encoder, decoder):
        super(MyFirstAE, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
    
    def forward(self, x):
        """
        Take a mini-batch as an input, encode it to the latent space and decode back to the original space
        x_out = decoder(encoder(x))
        :param x: torch.tensor, (MB, x_dim)
        :return: torch.tensor, (MB, x_dim)
        """
        # BEGIN SOLUTION (do not delete this comment!)

        x = self.encoder(x)
        x = self.decoder(x)

        # END Solution (do not delete this comment!)
        
        return x


# + [markdown] id="rp39DZHDAzfQ"
# #### It is high time to create encoder and decoder neural networks!
# Make hidden size of the network to be equal to `2`.
#
# **Hint.** You can use `nn.Sequential` to create your own archtectures.

# + id="gQ5Zuro3Aqmu"
encoder = lambda hid: nn.Sequential(
                            nn.Linear(30, 20),   
                            nn.LeakyReLU(inplace=True),
                            nn.Linear(20, 10),
                            nn.LeakyReLU(inplace=True),
                            nn.Linear(10, hid)
                        )  

decoder = lambda hid: nn.Sequential(
                            nn.Linear(hid, 10),
                            nn.LeakyReLU(inplace=True),
                            nn.Linear(10, 20),
                            nn.LeakyReLU(inplace=True),
                            nn.Linear(20, 30),
                        )

# + id="NOiJCm00A2t5"
device = 'cpu'

from torch.optim.lr_scheduler import StepLR

net = MyFirstAE(encoder(2), decoder(2))  
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.01, weight_decay=0.001)
scheduler = StepLR(optimizer, 30, gamma=0.5)
train_loader = torch_data.DataLoader(train_dset, batch_size=50, shuffle=True) 
val_loader = torch_data.DataLoader(val_dset, batch_size=200, shuffle=False) 


# + [markdown] id="hgOC1-iZA6ev"
# #### 7.4 Implement the missing parts of the `train` function

# + id="3SvYM1dxA5kI"
def train(epochs, net, criterion, optimizer, train_loader, val_loader,scheduler=None, verbose=True, save_dir=None):
    
    freq = max(epochs//20,1)
    net.to(device)
    
    for epoch in range(1, epochs+1):
        net.train()

        losses_train = []
        for X, _ in train_loader:
            
            ### BEGIN Solution

            # Perform one step of minibatch stochastic gradient descent
            reconstruction = net.forward(X)
            optimizer.zero_grad()
            loss = criterion(X, reconstruction)
            loss.backward()
            optimizer.step()
            losses_train.append(loss.item())

        # define NN evaluation, i.e. turn off dropouts, batchnorms, etc.
        net.eval()
        
        
        ## replace from loop
        losses_val = []
        for X, _ in val_loader:
            # Compute the validation loss

            with torch.no_grad():
                reconstruction = net.forward(X)
            loss = criterion(X, reconstruction)
            losses_val.append(loss.item())

            ### END Solution

        if scheduler is not None:
            scheduler.step()
        
        if verbose and epoch%freq==0:
            mean_val = sum(losses_val)/len(losses_val)
            mean_train = sum(losses_train)/len(losses_train)

            print('Epoch {}/{} || Loss:  Train {:.4f} | Validation {:.4f}'\
                  .format(epoch, epochs, mean_train, mean_val))

# + [markdown] id="IsrutdnJAasT"
# #### 7.5 Train your AE on breast cancer dataset. 
# Your goal is to get validation error <0.3.
#
# Some features that may help you to improve the performance:
# * `Dropout`
# * `Batchnorm`
# * lr scheduler
# * Batch size increase/decrease

# + colab={"base_uri": "https://localhost:8080/"} id="Gj9Bk-RQBHcD" outputId="63f79d3b-ff91-4238-e8f3-fe4a32fdcb73"
# for `MSE` loss get < 0.3 on validation, with AE "bottleneck" = 2
train(100, net, criterion, optimizer, train_loader, val_loader, scheduler) 

# + [markdown] id="9Tq4AMlDBCjW"
# #### 7.5 Let us take a look at the latent space. 
# Encode the whole dataset, using your AE, plot it in 2D and use colors to indicate objects of differrent classes

# + colab={"base_uri": "https://localhost:8080/", "height": 336} id="_DD8qANbBN1s" outputId="31562596-70df-4719-a1d0-7bf4ba625541"
 ### BEGIN Solution
    
plt.figure(figsize=(14, 5))
net.eval()
with torch.no_grad():
    enc = net.forward(torch.from_numpy(scaler.transform(cancer_dset['data'])).float()).detach().cpu()
plt.scatter(enc[:,0], enc[:,1], c=cancer_dset['target'], alpha=0.7);
plt.title('Latent space from the Autoencoder bottle neck, purple dots go for malignant samples. ');

 ### END Solution

# + [markdown] id="ufty_3qKBCwD"
# ### Task 8.  Autoencoder on kMNIST. (2 points)
#
#
# We will build a latent representation for `kMNIST` dataset by using our AE.
#
# 1. **(1 pt.)**  Train AE to get latent representation of the `kMNIST` dataset from `sklearn.datasets()`. Follow the instructions. Use `MSE` loss and obtain < $0.035$ on validation, with AE "bottleneck"  $\leq 40$; **(sub tasks 8.1 - 8.2)**
# 2.  **(1 pt.)** Plot 10 images and their reconstructions 2d. **(sub tasks 8.3)**

# + colab={"base_uri": "https://localhost:8080/", "height": 437, "referenced_widgets": ["9ea812fbae8b49a394df152e81fd359b", "d5177c89f9f64109809b98e3f7d12dcc", "ce378fbbbb4241abb1ccef1dadb77d04", "ff482e7b527e469bb6b56ebeb86f1576", "36ed9e979101449ba503116f3edf275e", "5945a58a17d24d5693703666efa8d714", "ce0259dd657e42d3988c92f5342c77a2", "e5058123f36344a0a6fec92a14751443", "611486a8f5844a6b84cad64a25eddfd7", "9ed16aa24f52411d904293d9c133958c", "15be64c126e1483f97cbb316cdd1f387", "627c9bca8a3c48d8a0b2567060bffd06", "b1bd9def1565419eb45a9ed7791612d6", "3757c84e012f4185a54d142311afaa70", "b14c5c5076e740348e751c58947f309d", "af21568ed3de4580beecc796f2c1574f", "5128c5524e7a43bda2ee14b3ecf90cf2", "412eeabcdd4d44c5b57a6fef085f913a", "f44e9c37697b45be88e25cfaa8756d63", "bce0850f0c9447a3919601d75485fc2b", "6d6d75dc2de249a7b03f3f5395619313", "5f54d80b86044766a7dfe3e49a71b46a", "94b87f82e52c4a27972ed4ba92503310", "5756dd971756457abac2a4a8e717047b", "3efa3908aa3c4eacb1671f176bdff33e", "ee76ea60574d43bbbd2cde6d98bd60f3", "d38d3449e1844015819d76f560c5718e", "d078c6bc68d543d38cf0e56253c6d3d5", "f05bda5aaec34aeb99edfe3e648926d3", "9f78f5369e6644f79de3dd05d20328b4", "8a8dcda652ef41cea608d8537a4ae845", "6cb3a11cf2c24fbc89e14e0c054c2404"]} id="mldP_RZZN7bm" outputId="8124e750-c7ad-45c2-c0a3-50399581e997"
from torchvision.datasets import KMNIST

data_train = KMNIST(train=True, root='./kmnist', download=True)
data_test = KMNIST(train=False, root='./kmnist', download=True)

# + [markdown] id="KZsmSJ3vuQrd"
# #### 8.1 Prepare the data and necessary functions.

# + id="KoRlKg3vOZCW"
x_train = np.array(data_train.data)
y_train = np.array(data_train.targets)

x_test = np.array(data_test.data)
y_test = np.array(data_test.targets)


# + id="fM3bsc0wBWTH"
# Reshape the data and scale

from sklearn.preprocessing import MaxAbsScaler

scaler = MaxAbsScaler()
n_train, n_test = x_train.shape[0], x_test.shape[0]
scaler.fit(x_train.reshape((n_train, -1)))
x_train = scaler.transform(x_train.reshape((n_train, -1))).reshape(n_train, 1, 28, 28)
x_test = scaler.transform(x_test.reshape((n_test, -1))).reshape(n_test, 1, 28, 28)

# + colab={"base_uri": "https://localhost:8080/", "height": 125} id="Tz2892txBYJk" outputId="6473c845-2646-4555-e82f-8a9ccae48a1a"
fig, ax = plt.subplots(ncols=10, figsize=(20, 5))
for i in range(10):
    ax[i].imshow(scaler.inverse_transform(x_train[i].reshape(1,-1)).reshape(28,28));
    ax[i].axis('off')


# + id="b5sHmYVxBeCV"
# BEGIN SOLUTION (do not delete this comment!)
class kMNISTData(torch_data.Dataset):
    def __init__(self, X, y):
        super(kMNISTData, self).__init__()
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.float32)
    
    def __len__(self):
        return  len(self.X)
    
    def __getitem__(self, idx):
        return self.X[idx].to('cuda'), self.y[idx].to('cuda')
    
# END Solution (do not delete this comment!)


# + id="lfla6EKiBgNy"
train_kmnist = kMNISTData(x_train, y_train) 
test_kmnist = kMNISTData(x_test, y_test) 


# + [markdown] id="J1jIdmh9uI-r"
# #### 8.2 Create encoder and decoder network for kMNIST. 
# You can either use convolutions or flatten the images and use linear layers. You can choose hidden size (not larger than 40) and any architecture you like. 

# + id="uvECUsVcBkmB"
# BEGIN SOLUTION (do not delete this comment!)

class Reshape(nn.Module):
    def __init__(self, *shape):
        super(Reshape, self).__init__()
        self.shape = shape

    def forward(self, x):
        return x.view(self.shape)
    
encoder = lambda hid: nn.Sequential(
                            nn.Conv2d(1, 12, 2, 2),
                            nn.BatchNorm2d(12),
                            nn.LeakyReLU(),
                            nn.Conv2d(12, 12, 3, 1, 1),
                            nn.BatchNorm2d(12),
                            nn.LeakyReLU(),
                            nn.Conv2d(12, 12, 2, 2),
                            nn.BatchNorm2d(12),
                            nn.LeakyReLU(),
                            nn.Conv2d(12, 6, 3, 1, 1),
                            nn.BatchNorm2d(6),
                            nn.Flatten(),
                            nn.Linear(6*7*7, hid)
                        ) 

decoder = lambda hid: nn.Sequential(
                            nn.Linear(hid, 6 * 7 * 7),
                            Reshape(-1, 6, 7, 7),
                            nn.BatchNorm2d(6),
                            nn.LeakyReLU(),
                            nn.Conv2d(6, 12, 3, 1, 1),
                            nn.BatchNorm2d(12),
                            nn.ConvTranspose2d(12, 12, 2, 2),
                            nn.BatchNorm2d(12),
                            nn.LeakyReLU(),
                            nn.Conv2d(12, 6, 3, 1, 1),
                            nn.BatchNorm2d(6),
                            nn.LeakyReLU(),
                            nn.ConvTranspose2d(6, 6, 2, 2),
                            nn.BatchNorm2d(6),
                            nn.LeakyReLU(),
                            nn.Conv2d(6, 1, 3, 1, 1),
                        ) 

# END Solution (do not delete this comment!)

# + id="Id-iNSswBpe9"
# BEGIN SOLUTION (do not delete this comment!)

device = 'cuda'

epochs = 25
net = MyFirstAE(encoder(40), decoder(40)) 
criterion = nn.MSELoss()

optimizer = torch.optim.Adam(net.parameters(), lr=0.01, weight_decay=0.001)
scheduler = StepLR(optimizer, 10, gamma=0.2)

train_loader = torch_data.DataLoader(train_kmnist, batch_size=100, shuffle=True) 
val_loader = torch_data.DataLoader(test_kmnist, batch_size=500, shuffle=False) 

# END Solution (do not delete this comment!)

# + colab={"base_uri": "https://localhost:8080/"} id="ytDZI0spBsFl" outputId="62697d23-46b4-4a44-a050-ad5199e96e81"
train(epochs, net, criterion, optimizer, train_loader, val_loader, scheduler)

# + [markdown] id="cXCR-eKBBuRI"
# #### 8.3 Plot any 10 images and their reconstructions.

# + colab={"base_uri": "https://localhost:8080/", "height": 290} id="ggGOxCc4Bvm5" outputId="fbfbc06a-2622-45ac-fd8e-3f99be6b581f"
# BEGIN SOLUTION (do not delete this comment!)

fig, ax = plt.subplots(ncols=10, nrows=2, figsize=(20, 5))
for i in range(10):
    im = train_kmnist[i][0]
    rec = net.forward(im.reshape(1,1,28,28)).detach().cpu().numpy()
    ax[0, i].imshow(scaler.inverse_transform(im.cpu().reshape(1,-1)).reshape(28,28));
    ax[1, i].imshow(scaler.inverse_transform(rec.reshape(1,-1)).reshape(28,28))
    ax[0, i].set_title('original')
    ax[1, i].set_title('reconstruction')
    ax[0, i].axis('off')
    ax[1, i].axis('off')
        
# END Solution (do not delete this comment!)

# + [markdown] id="1seXNwq3KoYM"
# ## Task 9. Convolutional NN (4 points)
#
#
# In this task, you will need to answer two questions and train a convolution neural network for a task of sound classification. 
#
# - **(1 pt.)** Debug the given convolutional neural network and explain what's wrong with it and how to fix it. You will need to identify at least 4 problems; **(sub-tasks 9.1)**
#
# - **(1 pt.)** Compute manually outputs of each layer, often when we build a neural network we need to know the output sizes of a layer before we add the next on; **(sub-tasks 9.2)**
#
# - **(2 pt.)** Build your own convolutional NN and train it for the task of sound classification. Your goal is to achieve maximum quality > 70%  1pt and > 90% 2pt. **(sub-tasks 9.3 - 9.6)**

# + [markdown] id="4fCPSsn3K22j"
# #### 9.1 Debug this convolutional neural network and write down proposed fixes. Ther are at least four fixes that can be applied. Explain your answers.

# + id="jQDXJDhFLI6a"
# assuming input shape [batch, 3, 32, 32]
cnn = nn.Sequential(
    nn.Conv2d(in_channels=1, out_channels=512, kernel_size=(3,3)), # 30
    nn.Conv2d(in_channels=512, out_channels=128, kernel_size=(3,3)), # 28
    nn.Conv2d(in_channels=128, out_channels=10, kernel_size=(3,3)), # 26
    nn.ReLU(),
    nn.MaxPool2d((1,1)),
    nn.Conv2d(in_channels=10, out_channels=3, kernel_size=(10,10)), # 17
    nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(10,10)), # 8
    nn.MaxPool2d((15,15)),
    nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(10,10)),
    nn.Softmax(),
    Flatten(),
    nn.Linear(64, 256),
    nn.Softmax(),
    nn.Linear(256, 10),
    nn.Sigmoid(),
    nn.Dropout(0.5)
    
)

# + id="0qomga4ALpdi"
# BEGIN SOLUTION (do not delete this comment!)

# Your answers:

# 1. Max poling with kernel 1 doesn't do anything
# 2. After 5-th convolution spatial size equal 8, but used maxpooling 
#    with kernel size 15; also problems after with sizes: if suppose 6-th conv produce
#    spatial size 1 linear layer should take 128 size instead of 64 
# 3. Softmax shouldn't be used not as last layer (same problems as sigmoid have);
#    for training purposes better use not softmax with nll loss, but use crossentropy
#    which is more numerically stable and more computationally efficient
# 4. Sigmoid bad choice because it is saturates fast; also better use bce with
#    logits function then place sigmoid as last layer with bce loss
# 5. Dropout should be inserted after linear layers except last layer (not before output)
# 6. Shortucts may increase quality, however network not very deep
# 7. No normalization layers which can increase stability
# 8. Fast dimension sizes changing commonly produce worse results than architectures with slowly changes sizes


# END Solution (do not delete this comment!)

# + [markdown] id="uh8BusFMQm-2"
# #### 9.2 Convolutional warm up, compute manually outpust shapes for each layer.

# + id="iGvJl1LxQoKw"
# Compute output shape for each and final layers wihout running the code.
# input size x = [8, 1 , 300, 303].

conv1 = Conv2d(in_chаnnels=1, out_channels=16, kernel_size=(5, 5), рadding = 0, stride=2)
conv2 = Conv2d(in_channels=16, out_chаnnels=16, kernel_size=(3, 3), рadding = 2, stride=1)
conv3 = Conv2d(in_channels=16, out_chаnnels=16, kernel_size=(5, 5), рadding =2, stride=2)
maxpool1 = MаxPool2d((2, 2))
cnn = nn.Sequential(conv1, conv2, conv3, maxpool1)


# + id="2ly8YywpLMKh"
# BEGIN SOLUTION (do not delete this comment!)

# example:

# conv1 
# output_h = (300 - 5+0) /2 +1 = 148
# output_w = (303 - 5+0) /2 +1 = 150

# Continue for all the layers:

### BEGIN Souluion

# conv2 
# output_h = (148 - 3 + 2 * 2) / 1 + 1 = 150
# output_w = (150 - 3 + 2 * 2) / 1 + 1 = 152

# conv3 
# output_h = (150 - 5 + 2 * 2) / 2 + 1 = 75
# output_w = (152 - 5 + 2 * 2) / 2 + 1 = 76

# maxpool1 = MaxPool2d((2, 2))
# output_h = 75 / 2 = 37
# output_w = 76 / 2 = 38

# final layer output = [8, 16, 37, 38] 

# END Solution (do not delete this comment!)

# + [markdown] id="1D1z6WEjfZwT"
# #### 9.3 Convolutional networks for sound classication
#
# - Now your task is to classify sounds using the convolutional network. You can use different network architectures. And your goal is to get the highest score possible.
#
# - First of all, we will preprocess audio into spectrograms, that you will be able to treat them as images.

# + colab={"base_uri": "https://localhost:8080/"} id="t7BsAPwYfv6X" outputId="d48bf5f3-53b7-4ed5-fd1a-836473edae21"
# imports

import os
import torch
import numpy as np
import torch.nn as nn
from torch import Tensor

# !pip install torchaudio
import torchaudio
from torchaudio import transforms
from IPython.display import Audio
import torch.nn.functional as F
from torch.utils.data import DataLoader,random_split,Dataset

import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, accuracy_score

# + colab={"base_uri": "https://localhost:8080/", "height": 66, "referenced_widgets": ["6d4a12e73a344069981447c9543053b9", "29c3287e1d1a476b89e8406f1c6f596a", "28af3e6e18824d7e9dcda0778c573883", "61216fae41b7448e8f724145e50d554e", "de7033ddcff84075be383e61ddb16c83", "c1996ab9c94d45c4aeea099f69441db3", "6b4ed0b03042452fa47a53f0313ddcba", "26aa4453751d4ced87119252e375ca11"]} id="lRCcakoVgFDK" outputId="1f9f1781-4e84-4cfb-dfe6-7dc549253328"
# Get the dataset

dataset = torchaudio.datasets.SPEECHCOMMANDS('./' , url = 'speech_commands_v0.02', 
                                       folder_in_archive= 'SpeechCommands', download = True)

# + [markdown] id="kUE5grAGj9SN"
# ### Let's look at the dataset.

# + colab={"base_uri": "https://localhost:8080/", "height": 282} id="85jqMauLhNER" outputId="076b5494-2ad8-41c1-d421-1a7caca8c449"
plt.figure()
plt.plot(dataset[0][0].t())

# + colab={"base_uri": "https://localhost:8080/", "height": 92} id="CYGRUGnghIYY" outputId="2ed60715-83b7-4903-d417-580d29ba17d9"
print('Label: ',dataset[11760][2])
Audio(np.array(dataset[11760][0].t()).reshape(-1), rate=16000)

# + [markdown] id="tMBl98nnkRcr"
# #### Actually, we could use really long sequences to classify our samples but it's better to work with them as spectrograms so we can use convolutional layers.

# + colab={"base_uri": "https://localhost:8080/", "height": 339} id="ZysQUugGkWPB" outputId="f56a9252-056e-4a17-be83-83d4a82a21e3"
specgram = torchaudio.transforms.Spectrogram(n_fft=200, normalized=True)(dataset[77][0])

print("Shape of spectrogram: {}".format(specgram.size()))

plt.figure(figsize=(10,5))
plt.imshow(specgram[0,:,:].numpy());
plt.colorbar()
plt.show()


# + colab={"base_uri": "https://localhost:8080/"} id="yylTFJjPo57s" outputId="7eac223f-3ec0-4c19-a027-9d8f973c2369"
# Some preprocessing routine
# Filter samples only with 16000 sampling rate
# Make labels dictionary

count = 0
wave = []
labels = []
labels_dict = {}

for i in range(0,105829):
    if dataset[i][0].shape == (1,16000):
        wave.append(dataset[i][0])
        labels.append(dataset[i][2])



# + colab={"base_uri": "https://localhost:8080/"} id="YD4Ds2BhqX4T" outputId="146e7799-5433-4f21-c313-d4c6d3512a31"
set_labels = list(set(labels))
labels_dict = {set_labels[i] :i for i in range(len(set_labels))}
labels_dict

# + [markdown] id="W5tkVQDTnMvD"
# #### 9.4 Your task right now is to implement a speech dataloader it will be almost the same as in the previous tasks.

# + id="kuGvJ4EDm90d"
transformation = torchaudio.transforms.Spectrogram(n_fft=200, normalized=True)

### BEGIN Solution

class SpeechDataLoader(Dataset):
    
    def __init__(self, data, labels, label_dict, transform=None):
        self.data = data
        self.labels = labels
        self.label_dict = label_dict
        self.transform = transform
            
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self,idx):
        waveform = self.data[idx]
        specgram = self.transform(waveform) 

        if self.labels[idx] in self.label_dict:
            label = self.label_dict[self.labels[idx]]

        return specgram, label
    
# END Solution (do not delete this comment!)


# + id="OarKlZWooQbS"
torch.manual_seed(0)

dataset= SpeechDataLoader(wave, labels, labels_dict, transformation)

traindata, testdata = random_split(dataset, [round(len(dataset)*.8), round(len(dataset)*.2)], )

train_loader = DataLoader(traindata, batch_size=100, shuffle=True)

val_loader = DataLoader(testdata, batch_size=100, shuffle=True)


# + [markdown] id="gQxaAGoBuLRR"
# #### 9.5 Your task is to build a convolutional neural network that  yields a high score.

# + id="ryDzP0l9s4Pi"
# BEGIN Solution  (do not delete this comment!)

class BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, relu=True, cropw=None, croph=None):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=2, stride=2),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels)
        )
        self.shortcut = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=2)
        self.cropw = cropw
        self.croph = croph
        self.relu = relu 

    def forward(self, x):
        out = self.backbone(x)
        if self.cropw is not None:
            x = x[:, :, :, :-self.cropw]
        if self.croph is not None:
            x = x[:, :, :-self.croph, :]
        out += self.shortcut(x)
        if self.relu:
            out = F.relu(out)
        return out
            
class NN2D(nn.Module):
    def __init__(self, num_class):
        super(NN2D,self).__init__()
        
        self.conv = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
        self.backbone = nn.Sequential(
            BasicBlock(32, 32, cropw=1, croph=1),
            BasicBlock(32, 64),
            BasicBlock(64, 128, croph=1),
            BasicBlock(128, 256),
            BasicBlock(256, 512, relu=False)
        )
        self.linear1 = nn.Linear(512, 128)
        self.dropout = nn.Dropout(inplace=True)
        self.linear2 = nn.Linear(128, num_class)
    
    def forward(self, x):
        
        out = self.conv(x)
        out = self.backbone(out)
        out = F.avg_pool2d(out, (3, 5))
        out = out.view(out.shape[0], -1)
        out = self.linear1(out)
        out = self.dropout(out)
        out = self.linear2(out)
        
        return out 
    
# END Solution (do not delete this comment!)


# + id="_Ev8ZHfquAMH"
# BEGIN Solution  (do not delete this comment!)

from torch.optim import Adam
from torch.nn.functional import cross_entropy

net = NN2D(len(set_labels)) 

num_epochs = 10
criterion = cross_entropy 
optimizer = Adam(net.parameters(), lr=0.001, weight_decay=0.001) 
scheduler = StepLR(optimizer, 7, gamma=0.2)

# END Solution (do not delete this comment!)

# + [markdown] id="NpgUGDXJvIhk"
# #### 9.6 Almost there, now, we need to rewrite our training loop a little bit.

# + id="k7BqVHuSvEhZ"
def train(epochs, net, criterion, optimizer, train_loader, val_loader,scheduler=None, verbose=True, device='cpu'):
    net.to(device)
    freq = max(epochs//15,1)
 
    for epoch in range(1, epochs+1):
        net.train()

        losses_train = []
        for X, target in train_loader:

            X, target = X.to(device), target.to(device)
            
            ### BEGIN Solution  (do not delete this comment!)
            
            # Perform one step of minibatch stochastic gradient descent
            
            predict = net.forward(X)
            optimizer.zero_grad()
            loss = criterion(predict, target)
            loss.backward()
            optimizer.step()
            losses_train.append(loss.item())
        
            # END Solution (do not delete this comment!)
            
        if scheduler is not None:
            scheduler.step()
        
        if verbose and epoch%freq==0:
            y_pred_val =  []
            y_true_val = []

            net.eval()
            # move from loop
            losses_val = []
            for X, target in val_loader:
                X, target = X.to(device), target.to(device)
                
                
                # BEGIN Solution  (do not delete this comment!)
                
                # Compute the validation loss
                
                with torch.no_grad():
                    target_hat_val = net.forward(X)
                
                loss = criterion(target_hat_val, target)
                losses_val.append(loss.item())
                
                # END Solution (do not delete this comment!)
                
                y_pred_val.extend(target_hat_val.argmax(1).tolist())
                y_true_val.extend(target.tolist())

            mean_val = sum(losses_val)/len(losses_val)
            mean_train = sum(losses_train)/len(losses_train)

            print('Val epoch {}'.format(epoch), \
              ', Loss : {:.3}'.format(mean_train), \
              ', Accuracy on test: {:.3}'.format(accuracy_score(y_true_val, y_pred_val)) )


# + colab={"base_uri": "https://localhost:8080/"} id="VjHENdc_uqc7" outputId="92400966-8348-46a9-f907-7f61234c963d"
train(num_epochs, net, criterion, optimizer, train_loader, val_loader, scheduler, device=0)
