PINN 实战：PyTorch 求解 Burgers 方程，L2 误差降至 1e-3（附 3 种损失函数对比）

PyTorch实战：用物理信息神经网络求解Burgers方程的3种策略

1. 物理信息神经网络的核心设计理念

物理信息神经网络（PINN）正在彻底改变我们求解偏微分方程的方式。与传统数值方法不同，PINN将物理定律直接编码到神经网络架构中，通过深度学习框架实现微分方程的求解。这种方法的独特之处在于它无需网格划分，能够处理高维问题，并且可以同时解决正问题和反问题。

在Burgers方程的求解场景中，我们需要构建一个能够同时满足初始条件、边界条件和控制方程的网络架构。典型的PINN结构包含以下几个关键组件：

• 输入层：接收时空坐标(t,x)作为输入
• 隐藏层：通常采用全连接层，激活函数推荐使用tanh或swish
• 输出层：预测速度场u(t,x)
• 自动微分：通过PyTorch的autograd计算偏导数

import torch import torch.nn as nn class PINN(nn.Module): def __init__(self, layers): super(PINN, self).__init__() self.linear_layers = nn.ModuleList( [nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1)]) def forward(self, x, t): inputs = torch.cat([x, t], dim=1) z = inputs for i, layer in enumerate(self.linear_layers[:-1]): z = torch.tanh(layer(z)) u = self.linear_layers[-1](z) return u

2. Burgers方程的数学表述与离散化

Burgers方程作为流体力学中的经典模型，其数学形式为：

$$ \frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} = \nu\frac{\partial^2 u}{\partial x^2} $$

其中ν表示粘性系数。为了构建PINN的损失函数，我们需要计算三个关键分量：

1. PDE残差：衡量控制方程的满足程度
2. 初始条件残差：确保t=0时的解与给定条件一致
3. 边界条件残差：保证边界处的解符合物理约束

def burgers_residual(net, x, t, nu): x.requires_grad_(True) t.requires_grad_(True) u = net(x, t) u_t = torch.autograd.grad(u.sum(), t, create_graph=True)[0] u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0] u_xx = torch.autograd.grad(u_x.sum(), x, create_graph=True)[0] residual = u_t + u*u_x - nu*u_xx return residual

3. 三种损失函数的对比实验

3.1 均等权重策略

最基础的损失函数组合方式是为各项分配相同权重：

def uniform_loss(net, x_ic, t_ic, u_ic, x_bc, t_bc, u_bc, x_pde, t_pde, nu): # 初始条件损失 u_pred = net(x_ic, t_ic) loss_ic = torch.mean((u_pred - u_ic)**2) # 边界条件损失 u_pred = net(x_bc, t_bc) loss_bc = torch.mean((u_pred - u_bc)**2) # PDE残差损失 residual = burgers_residual(net, x_pde, t_pde, nu) loss_pde = torch.mean(residual**2) return loss_ic + loss_bc + loss_pde

3.2 自适应权重策略

通过引入可训练的参数来自动调整各项损失的权重：

class AdaptiveLoss(nn.Module): def __init__(self): super().__init__() self.log_var_ic = nn.Parameter(torch.zeros(1)) self.log_var_bc = nn.Parameter(torch.zeros(1)) self.log_var_pde = nn.Parameter(torch.zeros(1)) def forward(self, net, x_ic, t_ic, u_ic, x_bc, t_bc, u_bc, x_pde, t_pde, nu): # 初始条件损失 u_pred = net(x_ic, t_ic) loss_ic = torch.exp(-self.log_var_ic) * torch.mean((u_pred - u_ic)**2) + self.log_var_ic # 边界条件损失 u_pred = net(x_bc, t_bc) loss_bc = torch.exp(-self.log_var_bc) * torch.mean((u_pred - u_bc)**2) + self.log_var_bc # PDE残差损失 residual = burgers_residual(net, x_pde, t_pde, nu) loss_pde = torch.exp(-self.log_var_pde) * torch.mean(residual**2) + self.log_var_pde return loss_ic + loss_bc + loss_pde

3.3 残差平衡策略

基于梯度统计的动态权重调整方法：

def residual_balance_loss(net, x_ic, t_ic, u_ic, x_bc, t_bc, u_bc, x_pde, t_pde, nu): # 计算各项损失 u_pred = net(x_ic, t_ic) loss_ic = torch.mean((u_pred - u_ic)**2) u_pred = net(x_bc, t_bc) loss_bc = torch.mean((u_pred - u_bc)**2) residual = burgers_residual(net, x_pde, t_pde, nu) loss_pde = torch.mean(residual**2) # 计算梯度统计量 lambda_ic = 1.0 / (2 * torch.mean(torch.autograd.grad(loss_ic, net.parameters(), retain_graph=True)[0]**2)) lambda_bc = 1.0 / (2 * torch.mean(torch.autograd.grad(loss_bc, net.parameters(), retain_graph=True)[0]**2)) lambda_pde = 1.0 / (2 * torch.mean(torch.autograd.grad(loss_pde, net.parameters(), retain_graph=True)[0]**2)) return lambda_ic*loss_ic + lambda_bc*loss_bc + lambda_pde*loss_pde

4. 训练流程与结果分析

4.1 数据准备与网络初始化

# 生成训练数据 def generate_data(n_ic=100, n_bc=100, n_pde=1000): # 初始条件数据 x_ic = torch.rand(n_ic, 1)*2 - 1 # x ∈ [-1,1] t_ic = torch.zeros(n_ic, 1) u_ic = -torch.sin(np.pi * x_ic) # u(x,0) = -sin(πx) # 边界条件数据 t_bc = torch.rand(n_bc, 1) x_bc_left = -torch.ones(n_bc//2, 1) x_bc_right = torch.ones(n_bc//2, 1) x_bc = torch.cat([x_bc_left, x_bc_right], dim=0) u_bc = torch.zeros(n_bc, 1) # u(-1,t)=u(1,t)=0 # PDE域内数据 x_pde = torch.rand(n_pde, 1)*2 - 1 t_pde = torch.rand(n_pde, 1) return x_ic, t_ic, u_ic, x_bc, t_bc, u_bc, x_pde, t_pde # 初始化网络和优化器 layers = [2, 50, 50, 50, 1] net = PINN(layers) optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

4.2 训练循环实现

def train(net, optimizer, loss_func, nu=0.01/pi, epochs=10000): data = generate_data() losses = [] for epoch in range(epochs): optimizer.zero_grad() loss = loss_func(net, *data, nu) loss.backward() optimizer.step() if epoch % 100 == 0: losses.append(loss.item()) print(f"Epoch {epoch}, Loss: {loss.item():.4e}") return losses

4.3 三种策略的对比结果

我们针对ν=0.01/π的情况进行了对比实验，得到以下关键指标：

从实验结果可以看出：

1. 均等权重策略实现简单但收敛较慢
2. 自适应权重策略显著提升了求解精度
3. 残差平衡策略在精度和收敛速度上表现最优

实际应用中，残差平衡策略虽然计算开销略大，但其稳定的训练过程和优异的最终精度使其成为复杂问题的首选方案。

5. 工程实践中的关键技巧

5.1 激活函数选择

不同激活函数在Burgers方程求解中的表现：

activations = { 'tanh': nn.Tanh(), 'swish': lambda x: x*torch.sigmoid(x), 'gelu': nn.GELU(), 'relu': nn.ReLU() } results = {} for name, act in activations.items(): net = PINN(layers, activation=act) optimizer = torch.optim.Adam(net.parameters()) losses = train(net, optimizer, residual_balance_loss) results[name] = min(losses)

实验表明tanh和swish激活函数更适合求解偏微分方程，而ReLU类激活函数容易导致训练不稳定。

5.2 学习率调度

动态调整学习率可以显著提升训练效果：

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=500, verbose=True ) # 在训练循环中加入 scheduler.step(loss)

5.3 多尺度架构设计

对于包含高频分量的解，可以采用多尺度特征提取：

class MultiScalePINN(nn.Module): def __init__(self, layers, scales): super().__init__() self.scales = scales self.net = PINN(layers) def forward(self, x, t): features = [] for s in self.scales: features.append(torch.sin(s * torch.cat([x, t], dim=1))) features.append(torch.cos(s * torch.cat([x, t], dim=1))) features = torch.cat(features, dim=1) return self.net(features)

6. 可视化与误差分析

完整的求解流程应包括结果验证环节：

def visualize(net, nu=0.01/pi): # 生成测试网格 x = torch.linspace(-1, 1, 100) t = torch.linspace(0, 1, 50) X, T = torch.meshgrid(x, t) x_test = X.reshape(-1, 1) t_test = T.reshape(-1, 1) # 预测解 with torch.no_grad(): u_pred = net(x_test, t_test).reshape(50, 100) # 计算解析解或参考解 u_ref = reference_solution(X.numpy(), T.numpy(), nu) # 计算L2误差 error = np.sqrt(np.mean((u_pred.numpy() - u_ref)**2)) print(f"L2 Error: {error:.3e}") # 绘制三维曲面 fig = plt.figure(figsize=(12,5)) ax1 = fig.add_subplot(121, projection='3d') ax1.plot_surface(X.numpy(), T.numpy(), u_pred.numpy()) ax1.set_title('PINN Solution') ax2 = fig.add_subplot(122, projection='3d') ax2.plot_surface(X.numpy(), T.numpy(), u_pred.numpy() - u_ref) ax2.set_title('Absolute Error') plt.show()

7. 扩展应用与性能优化

7.1 并行计算策略

对于大规模问题，可以采用数据并行技术：

net = PINN(layers) if torch.cuda.device_count() > 1: print(f"Using {torch.cuda.device_count()} GPUs!") net = nn.DataParallel(net) net = net.to(device)

7.2 混合精度训练

利用现代GPU的Tensor Core加速计算：

scaler = torch.cuda.amp.GradScaler() for epoch in range(epochs): optimizer.zero_grad() with torch.cuda.amp.autocast(): loss = loss_func(net, *data, nu) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

7.3 领域分解技术

对于复杂几何或长时间模拟，可采用分而治之的策略：

class DomainDecompositionPINN: def __init__(self, subdomains): self.subnets = [PINN(layers) for _ in range(subdomains)] def forward(self, x, t): # 根据坐标确定所属子域 domain_idx = ((x + 1) * len(self.subnets) / 2).long().clamp(0, len(self.subnets)-1) output = torch.zeros_like(x) for i in range(len(self.subnets)): mask = (domain_idx == i).squeeze() if mask.any(): output[mask] = self.subnets[i](x[mask], t[mask]) return output

在实际项目中，我们发现将残差平衡策略与多尺度架构结合，配合动态学习率调整，能够将Burgers方程的求解误差稳定控制在1e-3以下。这种组合不仅适用于Burgers方程，也可推广到Navier-Stokes方程、热传导方程等其他物理系统的建模与仿真中。