confidence.py

   1 #!/usr/bin/env python
   2
   3 import math
   4
   5 import torch, torchvision
   6
   7 from torch import nn
   8 from torch.nn import functional as F
   9
  10 ######################################################################
  11
  12 nb = 100
  13 delta = 0.35
  14 x = torch.empty(nb).uniform_(0.0, delta)
  15 x += x.new_full(x.size(), 0.5).bernoulli() * (1 - delta)
  16
  17 a = x * math.pi * 2 * 4
  18 b = x * math.pi * 2 * 3
  19 y = a.sin() + b
  20
  21 x = x.view(-1, 1)
  22 y = y.view(-1, 1)
  23
  24 ######################################################################
  25
  26 nh = 100
  27
  28 model = nn.Sequential(nn.Linear(1, nh), nn.ReLU(),
  29                       nn.Linear(nh, nh), nn.ReLU(),
  30                       nn.Linear(nh, 1))
  31
  32 criterion = nn.MSELoss()
  33 optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
  34
  35 for k in range(10000):
  36     loss = criterion(model(x), y)
  37     if (k+1)%100 == 0: print(k+1, loss.item())
  38     optimizer.zero_grad()
  39     loss.backward()
  40     optimizer.step()
  41
  42 ######################################################################
  43
  44 import matplotlib.pyplot as plt
  45
  46 fig, ax = plt.subplots()
  47 ax.scatter(x.numpy(), y.numpy())
  48
  49 u = torch.linspace(0, 1, 100).view(-1, 1)
  50 ax.plot(u.numpy(), model(u).detach().numpy(), color = 'red')
  51 plt.show()
  52
  53 ######################################################################