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

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/mlelarge/dataflowr/blob/master/CEA_EDF_INRIA/aleatoric_regression_colab.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="xnj4XU2aE2kO" colab_type="text"
# # Aleatoric regression

# + id="OaOS9uR1E2kR" colab_type="code" colab={}
# %matplotlib inline

import numpy as np
import torch
import math
import matplotlib.pyplot as plt
import imageio

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torchvision.utils import make_grid
from tqdm import tqdm, trange

import seaborn as sns


# + [markdown] id="2IaBfC4dE2kV" colab_type="text"
# ## Helper functions

# + id="TxRREopXE2kV" colab_type="code" colab={}
def normalize(image):
    """Takes a tensor of 3 dimensions (height, width, colors) and normalizes it's values
    to be between 0 and 1 so it's suitable for displaying as an image."""
    image = image.astype(np.float32)
    return (image - image.min()) / (image.max() - image.min() + 1e-5)


def display_images(images, titles=None, cols=5, interpolation=None, cmap="Greys_r"):
    """
    images: A list of images. I can be either:
        - A list of Numpy arrays. Each array represents an image.
        - A list of lists of Numpy arrays. In this case, the images in
          the inner lists are concatentated to make one image.
    """
    titles = titles or [""] * len(images)
    rows = math.ceil(len(images) / cols)
    height_ratio = 1.2 * (rows/cols) * (0.5 if type(images[0]) is not np.ndarray else 1)
    plt.figure(figsize=(15, 15 * height_ratio), dpi=200)
    i = 1
    for image, title in zip(images, titles):
        plt.subplot(rows, cols, i)
        plt.axis("off")
        # Is image a list? If so, merge them into one image.
        if type(image) is not np.ndarray:
            image = [normalize(g) for g in image]
            image = np.concatenate(image, axis=1)
        else:
            image = normalize(image)
        plt.title(title, fontsize=12)
        plt.imshow(image, cmap=cmap, interpolation=interpolation)
        i += 1

        
def show(img,cmap='viridis'):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest', cmap=cmap)


# + id="RT43ZKLoE2kZ" colab_type="code" colab={}
def uncertainity_estimate_2d(X, model, iters, l2, dropout_p=0.5, decay=1e-6):
    outputs = [model(X).cpu().data for i in range(iters)]
    outputs = torch.stack(outputs)
    y_mean = outputs.mean(dim=0)
    y_variance = outputs.var(dim=0)
    tau = l2 * (1. - dropout_p) / (2. * iters * decay)
    y_variance += (1. / tau)
    y_std = np.sqrt(y_variance)
    return y_mean, y_std



# + id="QQ36HKsSE2kc" colab_type="code" colab={}
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LOADER_KWARGS = {'num_workers': 1, 'pin_memory': True} if torch.cuda.is_available() else {}

# + id="eN6LHpaqE2ke" colab_type="code" colab={}
BATCH_SIZE = 128
TEST_BATCH_SIZE = 5

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        './mnist', train=True, download=True,
        transform=transforms.ToTensor()),
    batch_size=BATCH_SIZE, shuffle=True, **LOADER_KWARGS)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        './mnist', train=False, download=True,
        transform=transforms.ToTensor()),
    batch_size=TEST_BATCH_SIZE, shuffle=False, **LOADER_KWARGS)


TRAIN_SIZE = len(train_loader.dataset)
TEST_SIZE = len(test_loader.dataset)
NUM_BATCHES = len(train_loader)
NUM_TEST_BATCHES = len(test_loader)

CLASSES = 10
TRAIN_EPOCHS = 10 #20
SAMPLES = 2
TEST_SAMPLES = 10


# + id="lwkLkCbxE2kh" colab_type="code" colab={}
class AutoEncoderAleatoric(nn.Module):
    def __init__(self, p=0.5):
        self.dropout_p = p
        super(AutoEncoderAleatoric, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(True),
            nn.Linear(128, 64),
            nn.ReLU(True), 
            nn.Linear(64, 12), 
            nn.ReLU(True), 
            nn.Linear(12, 3))
        
        self.decoder = nn.Sequential(
            nn.Linear(3, 12),
            nn.ReLU(True),
            nn.Dropout(p=self.dropout_p),
            nn.Linear(12, 64),
            nn.ReLU(True),
            nn.Dropout(p=self.dropout_p),
            nn.Linear(64, 128),
            nn.ReLU(True), 
            nn.Dropout(p=self.dropout_p),
            nn.Linear(128, 128), 
            nn.Tanh())
        
        self.fc_pred =  nn.Linear(128, 28 * 28)
        self.fc_logvar =  nn.Linear(128, 28 * 28)


    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.encoder(x)
        x = self.decoder(x)
        pred = torch.tanh(self.fc_pred(x))
        logvar = F.softplus(self.fc_logvar(x))
        return pred, logvar


# + id="bSDp6emIE2kk" colab_type="code" colab={}
model = AutoEncoderAleatoric(p=0.1).to(DEVICE)


# + id="6cFcZu3FE2kn" colab_type="code" colab={}
def gaussian_log_loss(output, target, logvar):
    term1 = -0.5*torch.exp(-logvar)*(target - output)**2
    term2 = 0.5*logvar
    return (term1 + term2).sum()



# + id="hz3Oi_sfE2kp" colab_type="code" colab={}
def train_autoencoder(model, optimizer, criterion, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(tqdm(train_loader)):
#     for batch_idx, (data, target) in enumerate(train_loader):
        data = data.view(data.size(0), -1)
        data, target = data.to(DEVICE), target.to(DEVICE)
        model.zero_grad()
        out, logvar = model(data)
        loss = criterion(out, data, logvar)
        loss.backward()
        optimizer.step()
     
    print('epoch [{}], loss:{:.4f}'.format(epoch, loss.data.item()))



# + id="rcBrm00qE2ku" colab_type="code" colab={} outputId="962d2808-90b2-4391-875d-974f97f6291c"
TRAIN_EPOCHS = 5 #20
criterion = gaussian_log_loss
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(TRAIN_EPOCHS):
    train_autoencoder(model, optimizer, criterion, epoch)

# + id="FW7VCCtEE2k0" colab_type="code" colab={} outputId="695895ce-1ebc-4932-9573-6a1da3da5420"
mnist_sample, mnist_target = iter(test_loader).next()
mnist_sample = mnist_sample.to(DEVICE)
print(mnist_target)
# sns.set_style("dark")
show(make_grid(mnist_sample.cpu()))

# + id="ydLgSyrOE2k4" colab_type="code" colab={} outputId="e85dce15-6ad8-495a-f46e-cd64c34c8c3e"
out, logvar = model(mnist_sample)
out = out.detach().cpu()
out = out.view(-1, 1, 28,28)
logvar = logvar.detach().cpu()
logvar = logvar.view(-1, 1, 28,28)
print(out.size())
show(make_grid(out))

# + id="sQIvMh1uE2k7" colab_type="code" colab={} outputId="c04271e3-15ad-4746-bf2e-821594aeb395"
x = mnist_sample[1,:,:,:]
y_mean, y_std = out[1,:,:,:], torch.exp(logvar[1,:,:,:])

res_list = [x.cpu().view(28,28).numpy(), y_mean.view(28,28).numpy(), y_std.view(28,28).numpy()]
print(res_list[0].shape)
display_images(res_list, titles=['input','mean', 'std'], cols=len(res_list), interpolation=None, cmap="viridis")

# + [markdown] id="pYMhCjywE2k-" colab_type="text"
# ## Out of domain samples

# + id="4FD2Ynb0E2k_" colab_type="code" colab={} outputId="422ca501-e572-4043-8c61-debd31698454"
kmnist_loader = torch.utils.data.DataLoader(
    datasets.KMNIST('./kmnist', train=False, download=True, transform=transforms.ToTensor()), batch_size=5, shuffle=False)

# + id="Q4WPES8CE2lC" colab_type="code" colab={} outputId="c6326948-e7b4-4793-96a1-a3760af4d559"
kmnist_sample, kmnist_target = iter(kmnist_loader).next()
kmnist_sample = kmnist_sample.to(DEVICE)
print(kmnist_target)
# sns.set_style("dark")
show(make_grid(kmnist_sample.cpu()))

# + id="38OHYXZPE2lG" colab_type="code" colab={} outputId="ddccf943-ca8d-47dd-c14d-ff38e2bf4da0"
out, logvar = model(kmnist_sample)
out = out.detach().cpu()
out = out.view(-1, 1, 28,28)
logvar = logvar.detach().cpu()
logvar = logvar.view(-1, 1, 28,28)
print(out.size())
show(make_grid(out))

# + id="LucqxP4mE2lK" colab_type="code" colab={} outputId="95edb113-1f09-4d25-d8df-459ac285b4b5"
x = kmnist_sample[1,:,:,:]
y_mean, y_std = out[1,:,:,:], torch.exp(logvar[1,:,:,:])

res_list = [x.cpu().view(28,28).numpy(), y_mean.view(28,28).numpy(), y_std.view(28,28).numpy()]
print(res_list[0].shape)
display_images(res_list, titles=['input','mean', 'std'], cols=len(res_list), interpolation=None, cmap="viridis")

# + id="vtsAhXOTE2lR" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered