<jupyter_start><jupyter_text><jupyter_code>import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import StandardScaler
# 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
# visualize the data
x, y = generate_batch(1024)
plt.scatter(x, y, alpha=0.5);
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)
model = Model()
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
optimizer = torch.optim.Adam(model.parameters())
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()
plt.plot(losses);
# 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()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15324165029469547, 'fraction_numerical': 0.029469548133595286, 'mean_word_length': 3.506194690265487, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '353213', 'license': 'MIT', 'n_tokens_mistral': 1030, 'n_tokens_neox': 960, 'n_words': 242, 'path': 'PyTorch_Uncertainty_Estimation.ipynb', 'repo_name': 'abbaasalif/Uncertainity_estimator_pytorch', 'size': '2545'}