使用神经网络求解 ODE

Question

我想使用神经网络来解决这个常微分方程。 du/dt + 2u + t = 0，初始条件 u(0)=1，t 介于 0 到 2 之间。 我想用pytorch和自动微分法来求解这个方程。但我不知道如何在pytorch中计算du/dt。 我想将损失函数定义如下，并将其最小化以找到神经网络的最佳权重和偏差。 u_hat是用神经网络代替ODE的近似解。 R = du_hat/dt + 2*u_hat + t。 损失函数 = sum(R_i^2)。 损失函数是在 t = 0, 0.5, 1, 1.5, 2 点计算的 R_i 之和。

我不知道如何在pytorch中编写代码。

Answer 1

在下面的示例中，我首先使用标准求解器

scipy.integrate.solve_ivp

求解 ODE。该解决方案用于训练网络，因为它为每个

y

提供了一个目标

t

。网络将学习给定

t

和

y0

的参数，它将与参考

y

紧密匹配。

训练后，您可以向网络提供

t

和

y0

，网络将输出每个

y_hat

的估计解

t

。

请注意，这个示例有点小 - 您通常希望在未见过的样本（验证集）上评估模型，否则它可能只是记住训练数据而无法泛化到未见过的数据

t 

（尽管这对于您的用例来说可能不是问题）。

Net comprises 501 parameters
[epoch  100/ 500] loss: 0.00044
[epoch  200/ 500] loss: 0.00015
[epoch  300/ 500] loss: 0.00011
[epoch  400/ 500] loss: 0.00008
[epoch  500/ 500] loss: 0.00004

import torch
from torch import nn
from torch import optim

import numpy as np
import matplotlib.pyplot as plt

#Input data
n_samples = 100
t_array = np.linspace(0, 2, num=n_samples)
y0 = 1.0

#ODE function
# dy/dt + 2y + t = 0 --> dy/dt = -(2y + t)
def dy_dt(t, y):
    return -(2 * y + t)

#Solve using scipy
# The solution will be used to train the neural network
from scipy.integrate import solve_ivp
solved = solve_ivp(dy_dt, [0, 2], np.array([y_0]), t_eval=t_array)
solved_y = solved.y.ravel()

plt.plot(t_array, solved_y, color='cadetblue', linewidth=3, label='RK45 solver')
plt.xlabel('t')
plt.ylabel('y')
plt.gcf().set_size_inches(8, 3)

#
# Define the ODE net
#
class ODENet(nn.Module):
    def __init__(self, model_size=20, output_dim=1, activation=nn.ReLU):
        super().__init__()
        
        self.map_inputs = nn.Sequential(nn.Linear(2, model_size), activation())
        
        self.hidden_mapping = nn.Sequential(
            nn.Linear(model_size, model_size),
            activation()
        )
        
        self.output = nn.Linear(model_size, output_dim)
    
    def forward(self, x):
        # t, y0 = x[:, 0], x[:, 1]
        mapped_inputs = self.map_inputs(x)
        hidden = self.hidden_mapping(mapped_inputs)
        y_hat = self.output(hidden)
        return y_hat
    
print('Net comprises', sum([p.numel() for p in ODENet().parameters()]), 'parameters')

#Define the loss
# could alternatively pick from PyTorch's provided losses
def mse_loss(pred, target):
    return torch.mean((pred - target) ** 2)

#
# Create model
#
torch.manual_seed(0) #reproducible results

model = ODENet()
optimizer = optim.NAdam(model.parameters())

#Prepare the input data
# Convert to tensors
t_tensor = torch.Tensor(t_array).float().reshape(-1, 1)
y0_tensor = torch.Tensor([y0] * len(t_array)).float().reshape(-1, 1)
solved_y_tensor = torch.Tensor(solved_y).float()

# Combine inputs into a single matrix to make manpulation more compact
t_y0 = torch.cat([t_tensor, y0_tensor], dim=1)

#Scale the input features
# Will help the net's convergence, though not always abs necessary
t_y0 = (t_y0 - t_y0.mean(dim=0)) / (t_y0.std(dim=0) + 1e-10)

#
#Train model
#
n_epochs = 500
for epoch in range(n_epochs):
    model.train()
    
    y_hat = model(t_y0).flatten()

    #Loss, derivative, and step optimizer
    optimizer.zero_grad()
    loss = mse_loss(y_hat, solved_y_tensor)
    loss.backward()
    optimizer.step()
    
    #Print losses
    if ((epoch + 1) % 100) == 0 or (epoch + 1 == n_epochs):
        print(
            f'[epoch {epoch + 1:>4d}/{n_epochs:>4d}]',
            f'loss: {loss.item():>7.5f}'
        )

#Get the final predictions, and overlay onto the solver's solution
model.eval()
with torch.no_grad():
    predictions = model(t_y0)
    
plt.plot(t_array, predictions, color='sienna', linewidth=3, linestyle=':', label='ODENet')
plt.legend()

#Optional formatting
[plt.gca().spines[spine].set_visible(False) for spine in ['right', 'top']]
plt.gca().spines['bottom'].set_bounds(t_array.min(), t_array.max())
plt.gca().spines['left'].set_bounds(-0.75, 1)