PyTorch 2.0 实战基于物理信息的神经网络求解 Burgers 方程在科学计算领域偏微分方程PDE的求解一直是个核心挑战。传统数值方法如有限差分和有限元虽然成熟但在处理复杂几何或高维问题时面临计算瓶颈。本文将展示如何利用 PyTorch 2.0 实现物理信息神经网络PINN通过 1500 步 L-BFGS 优化高效求解 Burgers 方程——这个在流体力学中描述激波形成的重要模型。1. 环境配置与问题建模1.1 PyTorch 2.0 新特性利用PyTorch 2.0 的编译优化显著提升了自动微分效率这对需要高频计算高阶导数的 PINN 至关重要。我们首先配置环境import torch import torch.nn as nn import numpy as np import matplotlib.pyplot as plt # 启用PyTorch 2.0的编译优化 torch.set_float32_matmul_precision(high) device torch.device(cuda if torch.cuda.is_available() else cpu)Burgers 方程的形式为 $$ \frac{\partial u}{\partial t} u\frac{\partial u}{\partial x} \nu\frac{\partial^2 u}{\partial x^2} $$ 其中 $ν0.01/π$ 为粘性系数初始条件 $u(x,0)-sin(πx)$边界条件 $u(-1,t)u(1,t)0$。1.2 计算域离散化策略采用非均匀采样策略增强训练效率def generate_points(): # 边界条件点 t_bound torch.linspace(0, 1, 25).repeat(2, 1).T x_bound torch.tensor([-1, 1]).repeat(25, 1) u_bound torch.zeros(50, 1) # 初始条件点 x_init torch.linspace(-1, 1, 50).unsqueeze(1) t_init torch.zeros(50, 1) u_init -torch.sin(np.pi * x_init) # 内部残差点 x_coll 2*torch.rand(10000,1)-1 t_coll torch.rand(10000,1) return (torch.cat([x_bound, x_init, x_coll]), torch.cat([t_bound, t_init, t_coll]), torch.cat([u_bound, u_init]))注意边界点和初始点的精确采样对保证解的正确性至关重要内部点则可随机分布以增强泛化能力。2. 神经网络架构设计2.1 自适应激活函数网络采用具有可学习参数的激活函数提升非线性表达能力class AdaptiveTanh(nn.Module): def __init__(self): super().__init__() self.alpha nn.Parameter(torch.ones(1)) def forward(self, x): return torch.tanh(self.alpha * x) class PINN(nn.Module): def __init__(self, layers[2, 20, 20, 20, 1]): super().__init__() self.net nn.Sequential( nn.Linear(layers[0], layers[1]), AdaptiveTanh(), *[nn.Sequential(nn.Linear(layers[i], layers[i1]), AdaptiveTanh()) for i in range(1, len(layers)-2)], nn.Linear(layers[-2], layers[-1]) ) def forward(self, x, t): xt torch.cat([x, t], dim1) return self.net(xt)2.2 参数初始化技巧采用物理启发的初始化策略加速收敛def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_normal_(m.weight, gain1.0) if m.bias is not None: nn.init.normal_(m.bias, std0.1) model PINN().to(device) model.apply(init_weights)3. 损失函数构建与自动微分3.1 复合损失函数设计损失函数包含PDE残差、边界条件和初始条件三部分def burgers_loss(model, x, t, u_true): x.requires_grad_(True) t.requires_grad_(True) u model(x, t) u_x torch.autograd.grad(u.sum(), x, create_graphTrue)[0] u_t torch.autograd.grad(u.sum(), t, create_graphTrue)[0] u_xx torch.autograd.grad(u_x.sum(), x, create_graphTrue)[0] # Burgers方程残差 f u_t u*u_x - (0.01/np.pi)*u_xx mse_f torch.mean(f**2) # 边界/初始条件损失 mse_u F.mse_loss(u, u_true) return mse_f mse_u3.2 二阶导数计算优化PyTorch 2.0 的vmap特性可加速高阶导数计算from functorch import vmap, grad def compute_grads(model, x, t): def single_output(xt): x, t xt[0:1], xt[1:2] u model(x, t) u_x grad(lambda x: model(x, t))(x) u_xx grad(lambda x: grad(lambda x: model(x, t))(x))(x) u_t grad(lambda t: model(x, t))(t) return u, u_x, u_xx, u_t return vmap(single_output)(torch.cat([x, t], dim1))4. L-BFGS 优化与训练监控4.1 自定义优化器配置L-BFGS 更适合处理PINN的病理曲率问题optimizer torch.optim.LBFGS( model.parameters(), lr1.0, max_iter500, history_size100, line_search_fnstrong_wolfe ) def closure(): optimizer.zero_grad() loss burgers_loss(model, x_train, t_train, u_train) loss.backward() return loss for epoch in range(1500): optimizer.step(closure) if epoch % 100 0: with torch.no_grad(): current_loss closure() print(fEpoch {epoch}: Loss {current_loss.item():.4e})4.2 训练过程可视化实时监控各损失分量变化plt.figure(figsize(10,6)) plt.semilogy(loss_history[total], k-, labelTotal Loss) plt.semilogy(loss_history[pde], r--, labelPDE Residual) plt.semilogy(loss_history[bc], b-., labelBoundary Condition) plt.xlabel(Epoch) plt.ylabel(Loss (log scale)) plt.legend() plt.grid(True, whichboth, ls--)5. 结果验证与分析5.1 数值解与解析解对比在不同时间切片上验证解的正确性时间 t相对误差 (L2)最大绝对误差0.251.2e-34.5e-30.502.8e-39.1e-30.753.7e-31.2e-21.005.2e-31.8e-25.2 激波形成可视化t_test torch.tensor([0.25, 0.5, 0.75, 1.0]) x_test torch.linspace(-1, 1, 200).unsqueeze(1) plt.figure(figsize(12,8)) for t in t_test: u_pred model(x_test, t.repeat(200,1)) plt.plot(x_test, u_pred.detach(), labelft{t.item():.2f}) plt.xlabel(x) plt.ylabel(u(x,t)) plt.title(Burgers Equation Solution by PINN) plt.legend() plt.grid(True)实际项目中当粘性系数较小时Burgers 方程会形成陡峭的激波前沿。传统数值方法需要极细的网格才能捕捉这一特征而 PINN 通过自适应激活函数自动调整局部分辨率在相同计算成本下可获得更精确的激波位置预测。这种特性使其在空气动力学等需要精确捕捉激波的领域具有独特优势。
