torchdiffeq_neural_ode_quickstart_with_odeint_solver.py

1import torch
2import torch.nn as nn
3from torchdiffeq import odeint
4
5# Define the ODE system as a neural network
6class ODEFunc(nn.Module):
7    def __init__(self):
8        super(ODEFunc, self).__init__()
9        self.net = nn.Sequential(
10            nn.Linear(2, 50),
11            nn.Tanh(),
12            nn.Linear(50, 2),
13        )
14
15    def forward(self, t, y):
16        return self.net(y)
17
18# Initial conditions and time points
19y0 = torch.tensor([[2., 0.]])
20t = torch.linspace(0., 25., 100)
21
22# Initialize the function/model
23func = ODEFunc()
24
25# Solve the ODE: dy/dt = func(t, y)
26with torch.no_grad():
27    pred_y = odeint(func, y0, t)
28
29print(f"Output shape: {pred_y.shape}")
30# The result has shape [time_steps, batch_size, state_dimension]