class ODENet(nn.Module):
def __init__(self):
super(ODENet, self).__init__()
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
self.f = nn.Sequential(
nn.Linear(2, 50),
nn.GELU(),
nn.Linear(50, 75),
nn.GELU(),
nn.Linear(75, 1)
)
self.loss_fn = nn.MSELoss()
self.num_parameters = sum([np.prod(x.shape) for x in self.f.parameters()])
def forward(self, t, z):
if type(z) != torch.Tensor:
z = torch.Tensor(z)
z = z.to(device)
input_ = torch.cat([z, torch.Tensor([t]).to(device)])
dz = self.f(input_)
return dz
def forward_detached(self, t, z):
dz = self.forward(t, z).cpu().detach().numpy()
return dz
def get_z1(self, t0, t1, z0):
z0 = z0.cpu().numpy()
return torch.Tensor(solve_ivp(self.forward_detached, [t0, t1], z0, method='LSODA').y[:,-1]).to(device)
def loss(self, z1_pred, z1_true):
return self.loss_fn(z1_pred, z1_true)
def get_grad_cat(self):
grad = []
for layer in self.f:
if hasattr(layer, 'weight'):
grad.append(layer.weight.grad.flatten())
grad.append(layer.bias.grad.flatten())
return torch.cat(grad)
def aug(self, t, x):
self.zero_grad()
x = torch.Tensor(x).to(device)
z, a, _ = torch.split(x, [1, 1,x.shape[0]-2])
z.requires_grad = True
ft = self.forward(t, z)
ft.backward()
partial_f_partial_z = z.grad
partial_f_partial_theta = self.get_grad_cat()
da_dt = -a*partial_f_partial_z
partial_L_partial_theta = -a * partial_f_partial_theta
return torch.cat([ft, da_dt, partial_L_partial_theta]).cpu().detach().numpy()
def update_grad(self, t0, t1, z0, true_z1):
if type(z0) != torch.Tensor:
z0 = torch.Tensor(z0)
pred_z1 = self.get_z1(t0, t1, z0)
pred_z1.requires_grad = True
loss = self.loss(pred_z1, true_z1)
loss.backward()
partial_L_partial_z1 = pred_z1.grad
pred_z1.requires_grad = False
partial_L_partial_z1.requires_grad = False
x1 = torch.cat([pred_z1.cpu(), partial_L_partial_z1.cpu(), torch.zeros(self.num_parameters)])
x0 = torch.Tensor(solve_ivp(self.aug, [t1, t0], x1, method='LSODA').y[:,-1]).to(device)
_, _, grad = torch.split(x0, [1, 1,x0.shape[0]-2])
start_index=0
for layer in self.f:
if hasattr(layer, 'weight'):
n = layer.weight.flatten().shape[0]
layer.weight.grad = grad[start_index:start_index+n].reshape(layer.weight.grad.shape)
start_index += n
n = layer.bias.flatten().shape[0]
layer.bias.grad = grad[start_index:start_index+n].reshape(layer.bias.grad.shape)
start_index += n
return loss.item(), pred_z1
model = ODENet().to(device)
optimizer = torch.optim.Adam(model.parameters())
