confidence.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 ######################################################################
  16
  17 nb = 100
  18 delta = 0.35
  19 x = torch.empty(nb).uniform_(0.0, delta)
  20 x += x.new_full(x.size(), 0.5).bernoulli() * (1 - delta)
  21
  22 a = x * math.pi * 2 * 4
  23 b = x * math.pi * 2 * 3
  24 y = a.sin() + b
  25
  26 x = x.view(-1, 1)
  27 y = y.view(-1, 1)
  28
  29 ######################################################################
  30
  31 nh = 400
  32
  33 model = nn.Sequential(nn.Linear(1, nh), nn.ReLU(),
  34                       nn.Dropout(0.25),
  35                       nn.Linear(nh, nh), nn.ReLU(),
  36                       nn.Dropout(0.25),
  37                       nn.Linear(nh, 1))
  38
  39 model.train(True)
  40 criterion = nn.MSELoss()
  41 optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4)
  42
  43 for k in range(10000):
  44     loss = criterion(model(x), y)
  45     if (k+1)%100 == 0: print(k+1, loss.item())
  46     optimizer.zero_grad()
  47     loss.backward()
  48     optimizer.step()
  49
  50 ######################################################################
  51
  52 import matplotlib.pyplot as plt
  53
  54 fig, ax = plt.subplots()
  55
  56 u = torch.linspace(0, 1, 101)
  57 v = u.view(-1, 1).expand(-1, 25).reshape(-1, 1)
  58 v = model(v).reshape(101, -1)
  59 mean = v.mean(1)
  60 std = v.std(1)
  61
  62 ax.fill_between(u.numpy(), (mean-std).detach().numpy(), (mean+std).detach().numpy(), color = '#e0e0e0')
  63 ax.plot(u.numpy(), mean.detach().numpy(), color = 'red')
  64 ax.scatter(x.numpy(), y.numpy())
  65
  66 plt.show()
  67
  68 ######################################################################