import torch import torch.nn as nn import torch.optim as optim import matplotlib.pyplot as plt import numpy as np
步骤2： 定义 y(x) 的神经网络近似值
classODE_Net(nn.Module): def__init__(self, hidden_units=20): super(ODE_Net, self).__init__() self.layer1 = nn.Linear(1, hidden_units) self.layer2 = nn.Linear(hidden_units, hidden_units) self.layer3 = nn.Linear(hidden_units, 1) self.activation = nn.Tanh() defforward(self, x): out = self.activation(self.layer1(x)) out = self.activation(self.layer2(out)) out = self.layer3(out) return out
步骤3：计算 ODE 残差
defresidual(model, x): """ Compute the ODE residual: y''(x) - 3 = 0. """ # Enable gradients for x x.requires_grad_(True) y = model(x) # Compute first derivative: y'(x) dydx = torch.autograd.grad( y, x, grad_outputs=torch.ones_like(y), create_graph=True )[0] # Compute second derivative: y''(x) d2ydx2 = torch.autograd.grad( dydx, x, grad_outputs=torch.ones_like(dydx), create_graph=True )[0] # Compute the residual of the ODE: y''(x) - 3 res = d2ydx2 - 3.0 return res
步骤4：损失函数
defboundary_loss(model): """ Compute the loss associated with the boundary conditions: y(0)=1 and y(2)=5. """ # Boundary condition at x=0: y(0)=1 x0 = torch.tensor([[0.0]], device=device, requires_grad=True) y0 = model(x0) loss_bc0 = (y0 - 1.0)**2 # Boundary condition at x=2: y(2)=5 x2 = torch.tensor([[2.0]], device=device, requires_grad=True) y2 = model(x2) loss_bc2 = (y2 - 5.0)**2 return loss_bc0 + loss_bc2
步骤5：模型训练
# Initialize the model and optimizer model = ODE_Net(hidden_units=20).to(device) optimizer = optim.Adam(model.parameters(), lr=1e-3) num_epochs = 5000 # Generate interior points in the domain [0,2] N_interior = 50 x_interior = 2 * torch.rand(N_interior, 1, device=device) # uniformly distributed in [0,2] # Training loop for epoch in range(num_epochs): model.train() optimizer.zero_grad() # Compute the residual loss at interior points r_interior = residual(model, x_interior) loss_res = torch.mean(r_interior**2) # Compute the boundary condition loss loss_bc = boundary_loss(model) # Total loss is the sum of the residual and boundary losses loss = loss_res + loss_bc loss.backward() optimizer.step() if epoch % 500 == 0: print(f"Epoch {epoch}, Loss: {loss.item():.6e}") # Evaluate and compare the solution model.eval() x_test = torch.linspace(0, 2, 100, device=device).unsqueeze(1) y_pred = model(x_test).detach().cpu().numpy().flatten() x_test_np = x_test.cpu().numpy().flatten()
