# ---
# 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="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/BeatHubmann/legendary-giggle/blob/main/PyTorch_Uncertainty_Estimation.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="xaodMrOfUiIb"
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import StandardScaler


# + id="g-fgANGnU94U"
# data will come from a known distribution
# so we can check our answer, duh!
def generate_batch(batch_size=32):
  # x in (-5, +5)
  x = np.random.random(batch_size)*10 - 5

  # sd is a function of x
  sd = 0.05 + 0.1 * (x + 5)

  # target = mean + noise * sd
  y = np.cos(x) - 0.3 * x + np.random.randn(batch_size) * sd

  return x, y


# + id="QyAi_ahaWUT4" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="d72c186f-52e4-4d12-96d3-e22b8eb02c28"
# visualize the data
x, y = generate_batch(1024)
plt.scatter(x, y, alpha=0.5);


# + id="DllkLhrPaf7R"
class Model(nn.Module):
  def __init__(self):
    super().__init__()

    self.ann1 = nn.Sequential(
      nn.Linear(1, 10),
      nn.Tanh(),
      nn.Linear(10, 1),
    )
    self.ann2 = nn.Sequential(
      nn.Linear(1, 10),
      nn.Tanh(),
      nn.Linear(10, 1),
    )

  def forward(self, inputs):
    # returns (mean, log-variance)
    return self.ann1(inputs), self.ann2(inputs)


# + id="6onIwzs2aqFr"
model = Model()


# + id="BEDINlYhasAO"
def criterion(outputs, targets):
  mu = outputs[0]
  v = torch.exp(outputs[1])

  # coefficient term
  c = torch.log(torch.sqrt(2 * np.pi * v))

  # exponent term
  f = 0.5 / v * (targets - mu)**2

  # mean log-likelihood
  nll = torch.mean(c + f)

  return nll


# + id="KlamGAs1a1Bz"
optimizer = torch.optim.Adam(model.parameters())

# + id="GAylP2iDa5TY" colab={"base_uri": "https://localhost:8080/", "height": 102} outputId="28697fee-3ac8-42ec-ecdd-a3b3fd53f2e2"
n_epochs = 5000
batch_size = 128
losses = np.zeros(n_epochs)
for i in range(n_epochs):
  x, y = generate_batch(batch_size)

  # conver to torch tensor
  inputs = torch.from_numpy(x).float()
  targets = torch.from_numpy(y).float()

  # reshape data
  inputs, targets = inputs.view(-1, 1), targets.view(-1, 1)

  # zero grad
  optimizer.zero_grad()

  # forward pass
  outputs = model(inputs)
  loss = criterion(outputs, targets)

  # store loss
  losses[i] = loss.item()

  # print loss
  if i % 1000 == 0:
    print(i, losses[i])

  # optimize
  loss.backward()
  optimizer.step()

# + id="zAsm5ZHsbIPa" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="49a5b198-366b-4cfc-8bd5-134fef5057b6"
plt.plot(losses);

# + id="K5mHeCmZbIxV" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="521c4ce2-90b4-42a2-ca69-cb8ee0931808"
# plot the model predictions
x, y = generate_batch(1024)
plt.scatter(x, y, alpha=0.5)

# conver to torch tensor
inputs = torch.from_numpy(x).float()
targets = torch.from_numpy(y).float()

# reshape data
inputs, targets = inputs.view(-1, 1), targets.view(-1, 1)

with torch.no_grad():
  outputs = model(inputs)
  yhat = outputs[0].numpy().flatten()
  sd = np.exp(outputs[1].numpy().flatten() / 2) # since encoded variance by default

idx = np.argsort(x)
plt.plot(x[idx], yhat[idx], linewidth=3, color='red')
plt.fill_between(x[idx], yhat[idx] - sd[idx], yhat[idx] + sd[idx], color='red', alpha=0.3)
plt.show()

StarcoderdataJupyterScriptsDedupFiltered