# pasta

Sep 23rd, 2022 (edited)
917
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. import math
2. import numpy as np
3. from IPython.display import clear_output
4. from tqdm import tqdm_notebook as tqdm
5.
6. import matplotlib as mpl
7. import matplotlib.pyplot as plt
8. %matplotlib inline
9. import seaborn as sns
10. sns.color_palette("bright")
11. import matplotlib as mpl
12. import matplotlib.cm as cm
13.
14. import torch
15. from torch import Tensor
16. from torch import nn
17. from torch.nn  import functional as F
19.
20. use_cuda = torch.cuda.is_available()
21. def ode_solve(z0, t0, t1, f):
22.     """
23.    Простейший метод эволюции ОДУ - метод Эйлера
24.    """
25.     h_max = 0.05
26.     n_steps = math.ceil((abs(t1 - t0)/h_max).max().item())
27.
28.     h = (t1 - t0)/n_steps
29.     t = t0
30.     z = z0
31.
32.     for i_step in range(n_steps):
33.         z = z + h * f(z, t)
34.         t = t + h
35.     return z
36. class ODEF(nn.Module):
38.         """Compute f and a df/dz, a df/dp, a df/dt"""
39.         batch_size = z.shape[0]
40.
41.         out = self.forward(z, t)
42.
45.             (out,), (z, t) + tuple(self.parameters()), grad_outputs=(a),
46.             allow_unused=True, retain_graph=True
47.         )
48.         # метод grad автоматически суммирует градиенты для всех элементов батча,
49.         # надо expand их обратно
50.         if adfdp is not None:
53.         if adfdt is not None:
56.
57.     def flatten_parameters(self):
58.         p_shapes = []
59.         flat_parameters = []
60.         for p in self.parameters():
61.             p_shapes.append(p.size())
62.             flat_parameters.append(p.flatten())
65.     @staticmethod
66.     def forward(ctx, z0, t, flat_parameters, func):
67.         assert isinstance(func, ODEF)
68.         bs, *z_shape = z0.size()
69.         time_len = t.size(0)
70.
72.             z = torch.zeros(time_len, bs, *z_shape).to(z0)
73.             z[0] = z0
74.             for i_t in range(time_len - 1):
75.                 z0 = ode_solve(z0, t[i_t], t[i_t+1], func)
76.                 z[i_t+1] = z0
77.
78.         ctx.func = func
79.         ctx.save_for_backward(t, z.clone(), flat_parameters)
80.         return z
81.
82.     @staticmethod
83.     def backward(ctx, dLdz):
84.         """
85.        dLdz shape: time_len, batch_size, *z_shape
86.        """
87.         func = ctx.func
88.         t, z, flat_parameters = ctx.saved_tensors
89.         time_len, bs, *z_shape = z.size()
90.         n_dim = np.prod(z_shape)
91.         n_params = flat_parameters.size(0)
92.
93.         # Динамика аугментированной системы,
94.         # которую надо эволюционировать обратно во времени
95.         def augmented_dynamics(aug_z_i, t_i):
96.             """
97.            Тензоры здесь - это срезы по времени
98.            t_i - тензор с размерами: bs, 1
99.            aug_z_i - тензор с размерами: bs, n_dim*2 + n_params + 1
100.            """
101.             # игнорируем параметры и время
102.             z_i, a = aug_z_i[:, :n_dim], aug_z_i[:, n_dim:2*n_dim]
103.             # Unflatten z and a
104.             z_i = z_i.view(bs, *z_shape)
105.             a = a.view(bs, *z_shape)
109.
112.
119.
120.             # Flatten f and adfdz
121.             func_eval = func_eval.view(bs, n_dim)
124.
125.         dLdz = dLdz.view(time_len, bs, n_dim)  # flatten dLdz для удобства
127.             ## Создадим плейсхолдеры для возвращаемых градиентов
128.             # Распространенные назад сопряженные состояния,
129.             # которые надо поправить градиентами от наблюдений
132.             # В отличие от z и p, нужно вернуть градиенты для всех моментов времени
133.             adj_t = torch.zeros(time_len, bs, 1).to(dLdz)
134.
135.             for i_t in range(time_len-1, 0, -1):
136.                 z_i = z[i_t]
137.                 t_i = t[i_t]
138.                 f_i = func(z_i, t_i).view(bs, n_dim)
139.
140.                 # Рассчитаем прямые градиенты от наблюдений
141.                 dLdz_i = dLdz[i_t]
142.                 dLdt_i = torch.bmm(torch.transpose(dLdz_i.unsqueeze(-1), 1, 2),
143.                                    f_i.unsqueeze(-1))[:, 0]
144.
145.                 # Подправим ими сопряженные состояния
148.
149.                 # Упакуем аугментированные переменные в вектор
150.                 aug_z = torch.cat((
151.                     z_i.view(bs, n_dim),
154.                     dim=-1
155.                 )
156.
157.                 # Решим (эволюционируем) аугментированную систему назад во времени
158.                 aug_ans = ode_solve(aug_z, t_i, t[i_t-1], augmented_dynamics)
159.
160.                 # Распакуем переменные обратно из решенной системы
162.                 adj_p[:] += aug_ans[:, 2*n_dim:2*n_dim + n_params]
163.                 adj_t[i_t-1] = aug_ans[:, 2*n_dim + n_params:]
164.
165.                 del aug_z, aug_ans
166.
167.             ## Подправим сопряженное состояние в нулевой момент прямыми градиентами
168.             # Вычислим прямые градиенты
169.             dLdz_0 = dLdz[0]
170.             dLdt_0 = torch.bmm(torch.transpose(dLdz_0.unsqueeze(-1), 1, 2),
171.                                 f_i.unsqueeze(-1))[:, 0]
172.
173.             # Подправим
177. class NeuralODE(nn.Module):
178.     def __init__(self, func):
179.         super(NeuralODE, self).__init__()
180.         assert isinstance(func, ODEF)
181.         self.func = func
182.
183.     def forward(self, z0, t=Tensor([0., 1.]), return_whole_sequence=False):
184.         t = t.to(z0)
185.         z = ODEAdjoint.apply(z0, t, self.func.flatten_parameters(), self.func)
186.         if return_whole_sequence:
187.             return z
188.         else:
189.             return z[-1]