当前位置：首页 > news >正文

别再硬解方程了！用PyTorch搭建你的第一个物理信息神经网络（PINN），5分钟搞定一维热传导

news 2026/6/7 21:52:15

用PyTorch实现物理信息神经网络：5分钟攻克一维热传导方程

在传统工程计算中，解偏微分方程往往需要复杂的数值方法和繁琐的网格划分。但今天，我们将用PyTorch构建一个物理信息神经网络（PINN），让神经网络自己学会遵守物理定律。这种方法不仅代码简洁，还能避免传统数值方法中的稳定性问题。

1. 准备工作与环境配置

首先确保已安装PyTorch 1.8+版本。我们将使用以下核心组件：

import torch import torch.nn as nn import numpy as np import matplotlib.pyplot as plt

物理信息神经网络的关键在于将物理方程直接嵌入到神经网络的损失函数中。对于一维热传导问题，我们需要考虑：

数据损失：在已知数据点上网络输出与真实值的差异
物理损失：网络输出是否满足热传导方程

2. 构建神经网络模型

我们使用一个简单的全连接网络作为基础架构：

class HeatPINN(nn.Module): def __init__(self): super(HeatPINN, self).__init__() self.fc = nn.Sequential( nn.Linear(2, 20), # 输入(x,t) nn.Tanh(), nn.Linear(20, 20), nn.Tanh(), nn.Linear(20, 1) # 输出温度u ) def forward(self, x, t): inputs = torch.cat([x, t], dim=1) return self.fc(inputs)

这个网络有以下几个特点：

输入是空间坐标x和时间t的组合
使用Tanh激活函数保证输出平滑
网络深度适中，适合一维问题

3. 定义损失函数

热传导方程的物理约束是关键所在。一维热传导方程表示为：

∂u/∂t = α·∂²u/∂x²

在代码中实现物理约束：

def physics_loss(model, x, t, alpha=0.1): # 启用自动微分 x.requires_grad_(True) t.requires_grad_(True) u = model(x, t) # 计算一阶导数 du_dt = torch.autograd.grad(u.sum(), t, create_graph=True)[0] du_dx = torch.autograd.grad(u.sum(), x, create_graph=True)[0] # 计算二阶导数 d2u_dx2 = torch.autograd.grad(du_dx.sum(), x, create_graph=True)[0] # 物理残差 residual = du_dt - alpha * d2u_dx2 return torch.mean(residual**2)

4. 训练流程实现

完整的训练过程需要考虑数据采样和损失平衡：

def train(model, epochs=5000, lr=0.001): optimizer = torch.optim.Adam(model.parameters(), lr=lr) # 生成训练数据 x_data = torch.linspace(0, 1, 20).view(-1, 1) t_data = torch.linspace(0, 1, 20).view(-1, 1) for epoch in range(epochs): optimizer.zero_grad() # 随机采样内部点 x = torch.rand(100, 1, requires_grad=True) t = torch.rand(100, 1, requires_grad=True) # 计算各项损失 loss_physics = physics_loss(model, x, t) loss = loss_physics loss.backward() optimizer.step() if epoch % 500 == 0: print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

5. 结果可视化与验证

训练完成后，我们可以可视化预测结果：

def visualize(model): x_test = torch.linspace(0, 1, 100).view(-1, 1) t_test = torch.linspace(0, 1, 100).view(-1, 1) # 创建网格 X, T = torch.meshgrid(x_test.squeeze(), t_test.squeeze()) x_flat = X.reshape(-1, 1) t_flat = T.reshape(-1, 1) with torch.no_grad(): u_pred = model(x_flat, t_flat).reshape(100, 100) plt.figure(figsize=(10, 6)) plt.contourf(T.numpy(), X.numpy(), u_pred.numpy(), levels=20) plt.colorbar(label='Temperature') plt.xlabel('Time') plt.ylabel('Position') plt.title('PINN Solution for Heat Equation') plt.show()

6. 实战技巧与常见问题

在实际应用中，有几个关键点需要注意：

采样策略：边界点和内部点的采样比例会影响收敛
损失权重：数据损失和物理损失的相对权重需要调整
网络架构：更复杂的问题需要更深的网络

一个改进版的训练循环可能如下：

def improved_train(model, epochs=10000): optimizer = torch.optim.LBFGS(model.parameters()) def closure(): optimizer.zero_grad() # 边界条件采样 x_bc = torch.cat([ torch.zeros(50, 1), # x=0边界 torch.ones(50, 1) # x=1边界 ]) t_bc = torch.rand(100, 1) # 初始条件采样 x_ic = torch.rand(50, 1) t_ic = torch.zeros(50, 1) # 内部点采样 x_pde = torch.rand(500, 1) t_pde = torch.rand(500, 1) # 计算各项损失 loss_bc = ((model(x_bc, t_bc) - 0)**2).mean() # 假设边界温度为0 loss_ic = ((model(x_ic, t_ic) - torch.sin(np.pi*x_ic))**2).mean() loss_pde = physics_loss(model, x_pde, t_pde) loss = loss_bc + loss_ic + loss_pde loss.backward() return loss for epoch in range(epochs): optimizer.step(closure)