torchdiffeq_neural_ode_solver_minimal_example.py

1import torch
2import torch.nn as nn
3from torchdiffeq import odeint
4
5class ODEFunc(nn.Module):
6    def __init__(self):
7        super(ODEFunc, self).__init__()
8        self.net = nn.Sequential(
9            nn.Linear(2, 50),
10            nn.Tanh(),
11            nn.Linear(50, 2),
12        )
13
14    def forward(self, t, y):
15        return self.net(y)
16
17# Initial state
18y0 = torch.tensor([[2., 0.]])
19# Time points to evaluate
20t = torch.linspace(0., 25., 100)
21# Neural ODE model
22func = ODEFunc()
23
24# Solve the ODE
25# Method can be 'dopri5', 'adams', 'rk4', 'euler', etc.
26with torch.no_grad():
27    pred_y = odeint(func, y0, t, method='dopri5')
28
29print(f"Output shape: {pred_y.shape}")