Minor syntactic change.
[pytorch] / 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 ######################################################################