ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

物理信息神经网络(PINN)入门:PyTorch与TensorFlow框架对比与实现

物理信息神经网络(PINN)入门:PyTorch与TensorFlow框架对比与实现 PINN也就是物理信息神经网络Physics-Informed Neural Networks是这两年科学计算和深度学习交叉领域里关注度很高的方向。它的核心思路很简单把偏微分方程/常微分方程的残差作为损失项直接嵌入到神经网络训练过程中让网络不再只是“拟合数据”而是“遵守物理规律地去拟合数据”。这门技术最大的价值在于它能把传统数值求解 PDE 的思路和自动微分、GPU 并行、深度学习框架的能力结合起来在数据不足甚至没有数据的情况下也能求解偏微分方程。这篇内容是《PINN入门 30 讲》系列课程的第 11 讲重点放在 PINN 和深度学习框架的关系上。很多刚接触 PINN 的读者会卡在同一个地方到底选 PyTorch 还是 TensorFlow两个框架写 PINN 有什么区别自动微分怎么用损失函数里的 PDE 残差怎么算本讲就是回答这些问题的一次完整实操课。我们会从环境搭建开始分别用 PyTorch 和 TensorFlow 实现同一个 PINN 算例跑通训练并验证精度然后把训练过程中常见的坑全部列出来。适合刚了解 PINN 概念、准备开始写代码的读者。1. 核心能力速览项目属性说明项目类型PINN 入门系列课程第 11 讲框架选型与代码实现核心主题PINN物理信息神经网络、PyTorch、TensorFlow、自动微分主要功能用深度学习框架实现 PDE/ODE 求解、正向/反向问题建模推荐环境Python 3.9 及以上PyTorch 或 TensorFlow 二选一运行硬件CPU 可运行有 NVIDIA GPU 显存加速更佳显存占用以本讲一维 ODE 算例来说显存占用非常低一般显卡均可运行是否支持批量任务支持可通过批量构造采样点实现本讲验证的是单算例流程是否支持 API 服务本讲不涉及 API属于算法实现层适合场景科研入门、PINN 学习、框架选型评估、小规模 PDE 求解这一讲的目标不是讲解完整的 PINN 库而是把 PINN 最核心的计算流程“用两个主流深度学习框架各写一遍”。理解这一讲内容后再去读 DeepXDE、NVIDIA Modulus 这类 PINN 开源库会更容易上手。2. 适用场景与使用边界PINN 适合解决什么问题可以从正向和反向两个角度看。正向问题就是给定 PDE 和边界条件让网络学会输出解函数。比如热传导方程、波动方程、流体中的稳态或瞬态问题。反向问题则是已知部分观测数据利用 PINN 反推 PDE 中的未知参数。例如根据温度场离散点预测材料导热系数根据流速数据猜测粘性系数。这类问题的特点是传统数值方法往往需要网格生成和迭代求解器而 PINN 用统一框架就能把正反问题“一起解”尤其在数据稀疏的时候有优势。但 PINN 并不是万能的使用边界要说清楚。第一对于高维、多尺度、强非线性问题训练成本很高收敛稳定性也需要额外调优这时候传统有限元/有限体积法往往更可靠。第二PINN 输出是一个“近似解函数”不是网格节点上的离散精确解对误差有严格要求的工程场景需要与传统数值解交叉验证。第三如果问题本身没有明确的物理规律或者物理模型本身就不准PINN 并不会带来额外收益反而可能把模型错误固化到训练结果里。还有一个不能忽视的点是合规边界。PINN 属于数学建模和数值计算方法本身不涉及敏感内容。但实际使用时如果训练数据来自真实物理实验、工业检测或测绘数据要注意数据来源授权和隐私保护。使用 PyTorch 和 TensorFlow 开源框架时也要遵守各自的开源许可证生产环境发布前确认合规要求。3. PINN 环境准备PyTorch 与 TensorFlow 二选一搭建 PINN 环境前先理解为什么这两个框架都值得讲。PyTorch 采用动态图机制代码写起来更像 Python 原生逻辑调试方便科学计算社区相关代码很多。TensorFlow 的 Keras 接口非常简洁同时通过 GradientTape 也能完成自动微分在部分部署场景里有成熟方案。PINN 的本质是“神经网络 自动微分 损失函数组合”所以两个框架都满足需求区别更多在写代码的体验、生态组件和部署路径。环境准备遵循一个原则先把基础运行环境隔离好再安装深度学习框架最后用一段小代码验证自动微分可用。先创建一个独立的 conda 虚拟环境避免不同项目的依赖互相干扰conda create -n pinn python3.9 conda activate pinn然后根据是否有 NVIDIA GPU选择安装方式。PyTorch 和 TensorFlow 都支持 CPU 和 GPU 版本对于本讲算例CPU 就能轻松跑通。如果你机器上已经有 NVIDIA 驱动并安装了合适版本的 CUDA可以安装对应 GPU 版本加速训练如果没有独立显卡直接安装 CPU 版本即可不影响本讲理解 PINN 流程。安装时不要盲目复制网络上的命令尤其是 CUDA 版本必须和本机驱动匹配。最稳妥的做法是打开两个官网看当前稳定版本的安装命令。下面给出通用安装命令模板。4. PyTorch 安装与验证在 conda 环境中安装 PyTorchpip install torch如果希望安装 GPU 版本需要根据 PyTorch 官方安装页选择的 CUDA 版本用指定 index-url 安装。例如 CUDA 12.1 的安装命令通常类似pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121这里的 cu121 只是示例实际要按 PyTorch 官方页面提供的版本调整。安装完成后进入 Python 环境运行一段验证代码import torch print(PyTorch 版本:, torch.__version__) # 创建一个需要求梯度的张量 x torch.linspace(0, 2, 10, requires_gradTrue) y torch.sin(x) # 自动计算 dy/dx y.backward(torch.ones_like(y)) print(dy/dx:, x.grad)如果能够正常输出 dy/dx说明 PyTorch 和自动微分机制已经就绪。再验证一下 GPU 是否可用print(CUDA 是否可用:, torch.cuda.is_available())如果输出 True说明后续可以调用 GPU 训练如果输出 False也没关系本讲算例 CPU 就能完成。5. TensorFlow 安装与验证TensorFlow 的安装同样简单CPU 版本直接使用 pippip install tensorflowGPU 版本需要 CUDA、cuDNN 与 TensorFlow 版本匹配安装复杂度略高于 PyTorch。建议先跑通 CPU 版本再根据需求决定是否配置 GPU。安装后验证版本和自动微分能力import tensorflow as tf print(TensorFlow 版本:, tf.__version__) # GradientTape 自动微分 x tf.linspace(0.0, 2.0, 10) x tf.Variable(x) with tf.GradientTape() as tape: y tf.sin(x) dy_dx tape.gradient(y, x) print(dy/dx:, dy_dx.numpy())输出结果会是一个包含 cos(x) 数值的数组。GradientTape 是 TensorFlow 中实现 PINN 的关键工具后续代码里会重复用到。下面如果我们分别用 PyTorch 和 TensorFlow 实现同一个 PINN 算例你会看到它们在写法上的差别。6. 用 PyTorch 实现第一个 PINN一阶 ODE 算例现在进入核心实操。我们用一个非常经典的一阶常微分方程算例来演示dy/dx -yy(0) 1这个方程的真实解是 y e^{-x}。虽然问题简单但麻雀虽小五脏俱全它包含 PINN 的全部关键流程物理方程残差计算、边界条件处理、自动微分求导、损失加权训练。6.1 定义神经网络结构PINN 的神经网格不需要太复杂一般两三层的全连接网络加上 tanh 激活函数就能处理很多问题。这里我们定义一个 1-20-20-1 的 MLPimport torch import torch.nn as nn class PINN(nn.Module): def __init__(self): super(PINN, self).__init__() self.net nn.Sequential( nn.Linear(1, 20), nn.Tanh(), nn.Linear(20, 20), nn.Tanh(), nn.Linear(20, 1) ) def forward(self, x): return self.net(x)选择 tanh 而不是 ReLU是因为 PDE 求解要求输出函数具有一定的光滑性ReLU 输出存在不可导点会影响自动微分结果。6.2 构造损失函数PINN 的损失函数由 PDE 残差和边界条件两部分组成。PDE 残差就是让网络输出满足方程 dy/dx y 0 的逼近误差。在 PyTorch 中可以通过torch.autograd.grad自动计算网络输出对输入坐标的导数def pde_loss(model, x): x.requires_grad_(True) y model(x) dy_dx torch.autograd.grad( outputsy, inputsx, grad_outputstorch.ones_like(y), create_graphTrue, retain_graphTrue )[0] return torch.mean((dy_dx y) ** 2) def bc_loss(model, x_bc, y_bc): y_pred model(x_bc) return torch.mean((y_pred - y_bc) ** 2)在训练循环中x 的采样坐标在每次迭代可以直接重新生成也可以使用固定采样点。由于这里问题非常简单推荐使用固定采样点更容易复现# 在 [0, 2] 区间内采样内部点 x_train torch.linspace(0, 2, 100).reshape(-1, 1) # 边界条件y(0) 1 x_bc torch.zeros(1, 1) y_bc torch.ones(1, 1)6.3 训练循环model PINN() optimizer torch.optim.Adam(model.parameters(), lr0.01) for epoch in range(3000): optimizer.zero_grad() loss_pde pde_loss(model, x_train) loss_bc bc_loss(model, x_bc, y_bc) loss loss_pde loss_bc loss.backward() optimizer.step() if epoch % 500 0: print(fEpoch {epoch:5d}, Loss: {loss.item():.6e}, fPDE: {loss_pde.item():.6e}, BC: {loss_bc.item():.6e})看到 loss 逐渐下降说明模型开始学会了这个物理方程。训练完成后把预测结果和真实解 e^{-x} 对比一下import matplotlib.pyplot as plt import numpy as np x_test torch.linspace(0, 2, 100).reshape(-1, 1) y_pred model(x_test).detach().numpy() y_true np.exp(-x_test.numpy()) plt.plot(x_test.numpy(), y_true, labelTrue: e^{-x}) plt.plot(x_test.numpy(), y_pred, --, labelPINN Predict) plt.legend() plt.xlabel(x) plt.ylabel(y) plt.title(PINN solving dy/dx -y) plt.show()如果曲线贴合说明 PINN 已经正确学到了一阶 ODE 的解。这个流程虽然简单但后面无论面对热传导方程、流体方程还是更复杂的问题代码结构都不会变网络输出自动微分求导PDE 残差边界/初始条件损失加权优化。7. 用 TensorFlow 实现同样的 PINN同样的算例用 TensorFlow 再写一遍。这里最大的区别在于自动微分方式。PyTorch 使用torch.autograd.grad显式求导而 TensorFlow 使用tf.GradientTape作为上下文管理器来记录梯度。7.1 定义网络结构import tensorflow as tf model tf.keras.Sequential([ tf.keras.layers.Dense(20, activationtanh), tf.keras.layers.Dense(20, activationtanh), tf.keras.layers.Dense(1) ])7.2 定义复合损失函数def pinn_loss(): # 内部采样点 x_in tf.linspace(0.0, 2.0, 100) x_in tf.reshape(x_in, (-1, 1)) with tf.GradientTape() as tape: tape.watch(x_in) y model(x_in) dy_dx tape.gradient(y, x_in) # PDE 残差dy/dx y 0 pde_loss tf.reduce_mean(tf.square(dy_dx y)) # 边界条件y(0) 1 x_bc tf.zeros((1, 1)) y_bc model(x_bc) bc_loss tf.reduce_mean(tf.square(y_bc - 1.0)) return pde_loss bc_loss注意到一个细节在 GradientTape 中我们使用tape.watch(x_in)来追踪输入张量因为x_in不是模型中的可训练变量。如果没有这一行tape.gradient(y, x_in)会返回 None。7.3 训练循环optimizer tf.keras.optimizers.Adam(learning_rate0.01) for epoch in range(3000): with tf.GradientTape() as tape: loss pinn_loss() grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) if epoch % 500 0: print(fEpoch {epoch:5d}, Loss: {loss.numpy():.6e})TensorFlow 版和 PyTorch 版在结果上应该非常接近。你把两个实现放在一起对比就会发现核心数学完全一样区别只是 API 表达方式。这正好说明选 PyTorch 还是 TensorFlow不应该成为 PINN 入门的障碍关键是把自动微分和损失组合的逻辑理解透彻。8. 训练效果验证与性能观察PINN 训练完成后不能只看 loss 变低就认为成功还要做数值验证。上一节中我们将预测值和真实解析解画在一起做视觉比对这是第一步。更严谨一点可以计算 L2 相对误差# PyTorch 版验证 x_test torch.linspace(0, 2, 200).reshape(-1, 1) y_pred model(x_test).detach().numpy().flatten() y_true np.exp(-x_test.numpy()).flatten() l2_error np.linalg.norm(y_pred - y_true) / np.linalg.norm(y_true) print(fL2 相对误差: {l2_error:.6e})如果误差在 1e-3 到 1e-2 数量级说明训练是成功的。若误差偏大可以增加采样点数、加大网络宽度或增加训练轮次但要注意过拟合和训练时间成本。性能观察也是 PINN 实操中需要关注的点。在 CPU 上运行本讲算例训练时间通常在十几秒到几十秒之间显存占用很低。如果有 NVIDIA GPU训练会更快但对于这个小算例提升幅度不会太夸张因为网络规模和采样点数量都很小。训练时可以通过nvidia-smi查看显存占用观察程序是否真的使用了 GPUnvidia-smi如果看到 Python 进程占用显存说明 GPU 加速生效。对于后续更大规模的二维、三维 PDE 问题GPU 的作用会更明显所以环境准备阶段把 GPU 版本装好是有意义的。除了训练耗时还需要关注损失曲线形态。理想的损失曲线是平滑下降并趋于稳定。如果损失曲线震荡剧烈或者直接上升需要检查学习率是否过大、网络初始化是否合适、采样点分布是否合理。9. 常见问题与排查方法问题现象可能原因排查方式解决方案安装 PyTorch/torchvision 后 import 报错Python 版本不兼容或 pip 源异常检查 Python 版本和 pip 源升级到 Python 3.9切换官方 pip 源安装torch.cuda.is_available() 返回 FalsePyTorch 的 CUDA 版本与驱动不匹配运行nvidia-smi查看驱动版本按 PyTorch 官网选择匹配的 CUDA index-url 重装TensorFlow 版本过高导致 API 不兼容部分旧教程代码在新版本中变更查看版本迁移文档使用稳定版 API 或固定 TensorFlow 2.x 版本PyTorch 中backward()报错 element 0 of tensors does not require grad输入张量没有 requires_gradTrue或网络输出不含需要梯度的路径检查 x_train 是否设置 requires_grad检查模型参数是否可训练确认 x_train.requires_grad_() 已调用TensorFlow 中tape.gradient返回 None没有对输入张量执行 tape.watch确认在 pytorch 对应的是 requires_grad而 TensorFlow 中需要显式 tape.watch(x_in)在 GradientTape 中加入 tape.watch(x_in)损失不下降一直维持在很大值学习率过大/过小网络初始化不合理采样点未覆盖求解区域打印各损失分量观察 PDE 和 BC 哪个没有下降调小学习率使用 Xavier/He 初始化重采样训练点Loss 快速降到很低但预测曲线偏差大网络过拟合边界条件或 PDE 残差权重过小检查 loss_pde 和 loss_bc 比例直接绘制预测曲线调整损失权重让 PDE 和 BC 在损失中保持平衡比如 1:1 或按问题特征加权显存不足 (CUDA out of memory)网络规模或训练 batch 过大查看 nvidia-smi 中显存占用降低 batch size使用小网络或者切换到 CPU 先验证逻辑训练结果每次复现不一致随机初始化导致随机性固定随机种子设置 torch.manual_seed(42) 或 tf.random.set_seed(42)PINN 训练中还有一个非常典型的坑在 PyTorch 的autograd.grad中忘记设置create_graphTrue。如果我们需要高阶导数比如求解二阶扩散方程 u_xx必须让一阶导数的计算图被保留下来否则无法继续求二阶导。这个点是初学 PINN 最容易忽略的。# 二阶导数示例 def second_derivative(model, x): x.requires_grad_(True) u model(x) # 一阶导需要 create_graphTrue否则不能继续求二阶导 u_x torch.autograd.grad(u, x, grad_outputstorch.ones_like(u), create_graphTrue)[0] u_xx torch.autograd.grad(u_x, x, grad_outputstorch.ones_like(u_x), create_graphTrue)[0] return u_xx在 TensorFlow 中同样可以在嵌套的 GradientTape 中实现高阶导数def second_derivative_tf(model, x): with tf.GradientTape() as tape1: tape1.watch(x) with tf.GradientTape() as tape2: tape2.watch(x) u model(x) u_x tape2.gradient(u, x) u_xx tape1.gradient(u_x, x) return u_xx这两个例子已经足够覆盖绝大多数 PINN 编程需求了。后面遇到热传导方程、波动方程、扩散方程时你只需要把这里的导数阶数、残差项和数据项替换成对应方程形式即可。10. PINN 工程化最佳实践PINN 从“跑通一个算例”到“真正解决一个 PDE 问题”中间有不少工程化细节值得提前注意。这里分享几条对新手最实用的建议。第一条保持一个最小可用模板。本讲的 PyTorch 和 TensorFlow 代码就是很好的起点把网络定义、损失函数、训练循环拆成清晰函数或模块以后遇到新问题只需要修改方程残差的部分。不要每次从零写脚本那样容易引入低级错误。第二条训练前先做归一化。坐标 x 的取值范围如果过大神经网络很难收敛。一般建议把求解区域归一化到 [-1, 1] 或 [0, 1] 区间。对于物理量 y如果真值尺度过大或过小也可以做无量纲化处理。虽然本讲算例范围是 [0, 2]不归一化也能收敛但二维、三维问题会越来越敏感。第三条用损失分项日志替代单一总 loss。PINN 训练时把 PDE 残差、边界条件、数据拟合项分开打印。很多时候总 loss 在下降但边界条件这一项已经发散说明权重失衡。看到各分项数值变化才能判断问题出在哪。第四条固定随机种子。PINN 训练结果受网络初始化和采样点分布影响明显。固定随机种子能保证实验可复现这在写论文、调参数、比较算法时非常关键。第五条关于批量任务和参数扫描。如果你的目标是系统的参数研究比如改变方程系数、边界条件、网络宽度、学习率建议写一个简单的 shell 循环或 Python 脚本把采样点数量、训练轮次、网络层数作为参数传入并把训练日志和模型权重保存到独立目录。这样自动化跑实验会高效很多。# 伪代码示例批量参数扫描 for hidden_size in 10 20 50; do python train_pinn.py --hidden_size $hidden_size --epochs 3000 done第六条生产环境部署时留意模型导出和推理速度。PyTorch 可以导出 TorchScript 或 ONNX 格式TensorFlow 可以使用 SavedModel 格式都能做轻量化推理。不过 PINN 最主要的应用场景还是“求解验证”模型导出不是每篇文章都覆盖的重点但如有部署需求从第一步就使用干净的代码结构会带来方便。11. 总结与下一步本讲把 PINN 和两个主流深度学习框架之间的关系讲清楚了。我们分别用 PyTorch 和 TensorFlow 实现了同一个一阶 ODE 算例走通了“定义网络、计算 PDE 残差、施加边界条件、训练、可视化验证”的完整流程。两个框架的核心逻辑完全一致区别只在自动微分的 API 写法上。PyTorch 用torch.autograd.gradTensorFlow 用tf.GradientTape理解了这两者的用法后续无论换哪个框架写 PDE 求解思路都能平移。如果你是刚接触 PINN建议先在本讲代码上做两件事。第一件把这个一阶 ODE 改成二阶 ODE比如 u u 0u(0)1u(0)0体会一下高阶导数如何计算。第二件把训练采样方式从固定网格点改成每次迭代随机采样观察训练稳定性和收敛速度的变化。跑通这些改动后再进入热传导方程、二维 Poisson 方程等真正的 PDE 算例。下一讲我们会把 PINN 扩展到一个真正有用的 PDE 场景一维 Burgers 方程或一维热传导方程。到那时网络结构、损失权重、时间项离散这几个问题会一起出现建议收藏本讲代码作为基础模板。如果你手头正好有正在学习 PINN 的朋友这篇框架对比笔记可以直接转给他参考。
返回列表