1 import torch
2
3
4 nn = 50
5 model = torch.nn.Sequential()
6 model.add_module('layer_1', torch.nn.Linear(2,nn))
7 model.add_module('fun_1', torch.nn.Tanh())
8 model.add_module('layer_2', torch.nn.Linear(nn,nn))
9 model.add_module('fun_2', torch.nn.Tanh())
10 model.add_module('layer_3', torch.nn.Linear(nn,nn))
11 model.add_module('fun_3', torch.nn.Tanh())
12 model.add_module(f'layer_4', torch.nn.Linear(nn,1))
13
14
15
16
17 def fun(x1, x2):
18 return torch.sin(torch.pi*x1) * \
19 torch.sin(torch.pi*x2)
20
21 x1_a = -1.
22 x1_b = 1
23
24 x2_a = -1.
25 x2_b = 1.
26
27
28
29 optim = torch.optim.SGD(model.parameters(),
30 lr=1e-1, momentum=0.9)
31
32
33 ns = 20
34
35 nepochs = 50000
36
37 tol = 1e-4
38
39
40 n_val = 50
41 x1 = torch.linspace(x1_a, x1_b, steps=n_val)
42 x2 = torch.linspace(x2_a, x2_b, steps=n_val)
43 X1_val, X2_val = torch.meshgrid(x1, x2, indexing='ij')
44 X_val = torch.hstack((X1_val.reshape(n_val**2,1),
45 X2_val.reshape(n_val**2,1)))
46 Y_vest = fun(X1_val, X2_val).reshape(-1,1)
47
48 for epoch in range(nepochs):
49
50
51 X1 = (x1_b - x1_a) * torch.rand(ns**2, 1) + x1_a
52 X2 = (x2_b - x2_a) * torch.rand(ns**2, 1) + x2_a
53
54 X_train = torch.hstack((X1, X2))
55 Y_train = fun(X1, X2).reshape(-1,1)
56
57
58
59 Y_est = model(X_train)
60
61
62 loss = torch.mean((Y_est - Y_train)**2)
63
64 if (epoch % 100 == 0):
65 print(f'{epoch}: {loss.item():.4e}')
66
67
68 optim.zero_grad()
69 loss.backward()
70 optim.step()
71
72
73 if (epoch % 100 == 0):
74 Y_val = model(X_val)
75 loss_val = torch.mean((Y_val - Y_vest)**2)
76
77 print(f"\tloss_val = {loss_val.item():.4e}")
78
79
80 if (loss_val.item() < tol):
81 break