多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

PINN入门:PyTorch与TensorFlow框架选型及自动微分实现

PINN入门:PyTorch与TensorFlow框架选型及自动微分实现 PINN物理信息神经网络入门进行到第 11 讲主题正好落到框架上。前面的内容大多数在处理偏微分方程、损失函数、边界条件和物理约束这些概念但从这一讲开始所有想法都要在 PyTorch 或 TensorFlow 里变成能跑的代码。框架选型不是玄学它直接决定自动微分顺不顺手、断点好不好打、GPU 能不能用上、后期批量算例怎么组织。这篇文章围绕 PINN 最常用的两个后端——PyTorch 和 TensorFlow——把框架能力差异、环境安装、最小可运行算例和调参排查完整拆一遍。适合已经了解 PINN 基本思想、准备进入代码阶段的读者也适合正在犹豫先学哪个框架的人。1. PINN 的代码本质先看框架到底要承担什么很多人学 PINN 卡住不是数学看不懂而是不知道代码里哪些部分是框架在起作用。先把这个问题拆清楚后面选框架、装环境、调 bug 都会顺很多。1.1 一个 PINN 程序由三块组成第一块是神经网络本身。输入是坐标一维问题就是 x二维问题就是 x 和 y含时间就是 x、y、t输出是待求物理量比如位移、温度、速度或流函数。这一块用 PyTorch 和 TensorFlow 写都很简单本质上就是几层全连接网络加激活函数。第二块是自动微分。PINN 的损失函数里包含物理残差比如一维泊松方程 -u(x) f(x)你需要对网络输出 u 求关于输入 x 的二阶导数。这个导数不是用差分近似而是框架通过计算图自动求出来的精确且可反向传播。第三块是损失函数和优化器。损失一般由三部分拼接而成PDE 残差损失、边界条件损失、如果有观测数据再加数据损失。优化器用 Adam 起步就够后续可以换 L-BFGS 做精细收敛。1.2 自动微分是 PINN 和普通神经网络的分水岭普通分类或回归任务网络只做一次前向计算loss 是输出和标签的差。PINN 不一样loss 里要求导数而且二阶导很常见。这意味着框架必须支持对网络输出反复求导并且让导数继续参与梯度回传。这一条决定了框架选型。你在 PyTorch 里会用到torch.autograd.grad在 TensorFlow 里会用到tf.GradientTape。两者的使用思路相同但细节差异很大比如是否要create_graphTrue、是否要persistentTrue。这些细节就是 PINN 入门阶段最常见的报错来源。建议先把“自动微分求二阶导”这个动作练熟再去看复杂的 PINN 结构。这是整门课程里复用率最高的代码片段。2. PyTorch 和 TensorFlow 在 PINN 场景下的实际差异这两个框架都能写 PINN但写法风格、调试体验和部署路径不太一样。下面按 PINN 代码的实际组织方式对比。2.1 PyTorch显式 autograd高阶导数更顺手PyTorch 是动态图风格每次前向都会重新构建计算图。这个特点对 PINN 非常友好因为 PINN 的输入点经常是每次迭代重新采样图结构本来就在变。求一阶导u_x torch.autograd.grad( u, x, grad_outputstorch.ones_like(u), create_graphTrue )[0]grad_outputs是告诉框架“输出 u 的每一维权重都是 1”相当于对 sum(u) 求导。create_graphTrue是为了让导数本身也进入计算图这样你才能继续对u_x再求一次导得到u_xx。PyTorch 的调试体验接近普通 Python 脚本打印中间张量、打断点都自然。学术论文复现时PyTorch 代码可读性也更高所以你看到的新论文里 PINN 实现多半是 PyTorch。2.2 TensorFlowGradientTape 管理梯度Keras 集成方便TensorFlow 求导需要显式使用tf.GradientTape并且要先tape.watch(x)让框架关注输入张量with tf.GradientTape(persistentTrue) as tape: tape.watch(x) u model(x) u_x tape.gradient(u, x) u_xx tape.gradient(u_x, x)persistentTrue是必须的因为默认情况下 tape 只能调用一次gradient求完一阶导后 tape 资源就释放了。PINN 需要两阶甚至更高阶导数所以要把 tape 设为 persistent。TensorFlow 的优势是 Keras 高层 API 方便模型定义、训练循环、模型保存都有现成工具。如果之后要把训练好的 PINN 模型导出到线上推理环境TensorFlow 的 SavedModel 流程相对成熟。2.3 选型建议主力用一个另一个用来对照我的建议很直接课程学习阶段两个框架都装但主力只用一个。先用 PyTorch 把一个算例跑通再用 TensorFlow 复现同一个算例对比一次你就明白了框架差异到底在哪。如果只选一个从 PINN 入门和论文复现角度我倾向 PyTorch。原因不是 TensorFlow 不行而是 PyTorch 的自定义损失、自定义训练循环和调试点更贴近 PINN 的科研场景。TensorFlow 适合你已经熟悉 Keras 生态、后续要快速部署的场景。下表是常见关注点的对照关注点PyTorchTensorFlow一阶导torch.autograd.gradtape.gradient二阶导同一函数再次调用需create_graphTruepersistentTrue后再次调用网络定义nn.Modulenn.Sequentialtf.keras.Sequential训练循环手写 for 循环为主可手写也可model.fit调试打断点、打印张量都自然eager 模式可以图模式受限部署ONNX、TorchScriptSavedModel论文复现常见相对少3. 环境准备先把两套框架安装到能跑为止框架选好了下一步是安装。说实话PINN 入门阶段最常见的翻车点不是算法而是环境没配对Python 版本不对、CUDA 版本不匹配、GPU 版装成 CPU 版、两个框架装在同一个环境里互相冲突。3.1 用虚拟环境隔离避免两套框架互相干扰我建议用 Anaconda 或 Miniconda 管理 Python 环境。PyTorch 和 TensorFlow 的 CUDA 依赖并不完全相同强行装到一个环境里可能出现库冲突表现为某个框架导入失败或者 GPU 不可用。创建两个独立环境一个给 PyTorch一个给 TensorFlowconda create -n pinn-torch python3.10 -y conda create -n pinn-tf python3.10 -y用两个环境的好处是每个框架的依赖都能保持干净出问题也好排查。你只需要记清楚当前 activate 的是哪个环境。3.2 PyTorch 安装与 GPU 验证PyTorch 安装命令建议直接去官方 get-started 页面生成因为不同 CUDA 版本对应不同安装源。如果是纯学习 CPU 就够直接pip install torch需要 GPU 时官方给的命令一般是带--index-url的形式比如pip install torch --index-url https://download.pytorch.org/whl/cu121这个 CUDA 版本号必须和你机器上实际安装的驱动兼容。安装完成后用下面这段验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果torch.cuda.is_available()返回 False先别怀疑代码优先检查驱动、CUDA 版本和 PyTorch 安装源是否匹配。3.3 TensorFlow 安装与 GPU 验证TensorFlow 的 CPU 版安装很简单pip install tensorflow新版 TensorFlow 不再单独发布tensorflow-gpu包GPU 支持是否启用取决于本机 CUDA 和 cuDNN 版本是否匹配。不同小版本对 CUDA 的要求不一样建议安装前看一眼官方 pip 页面的版本说明。验证方式import tensorflow as tf print(tf.__version__) print(tf.config.list_physical_devices(GPU))如果 GPU 列表为空常见原因是 CUDA/cuDNN 版本不匹配或者驱动版本过旧。这时候先跑 CPU 版本学习完全没问题不要让 GPU 问题卡住课程进度。初学者先 CPU 跑通算例再折腾 GPU。PINN 小算例在 CPU 上几千步也就几十秒到几分钟算法能收敛比速度快更重要。4. 以 1D 泊松方程为例写一个 PyTorch 版 PINN光说概念没用直接看一个能跑的算例。这里选一维泊松方程-u(x) π² sin(πx)x ∈ [0, 1]边界条件 u(0) 0u(1) 0。这个方程的精确解是 u(x) sin(πx)方便最后验证误差。不用复杂方程因为目的是把 PINN 的代码骨架跑通。4.1 定义神经网络网络结构很简单输入 1 维输出 1 维中间三层全连接每层 20 个神经元激活函数用 tanh。tanh 在 PINN 里比 ReLU 常用因为 ReLU 的二阶导基本恒为零无法用来构造二阶方程残差。import torch import torch.nn as nn import torch.optim as optim class PINN(nn.Module): def __init__(self): super().__init__() self.net nn.Sequential( nn.Linear(1, 20), nn.Tanh(), nn.Linear(20, 20), nn.Tanh(), nn.Linear(20, 20), nn.Tanh(), nn.Linear(20, 1), ) def forward(self, x): return self.net(x)4.2 物理残差和损失函数核心函数是 PDE 残差。要注意两点一是对输入x要重新clone().requires_grad_(True)二是create_graphTrue必须开否则二阶导算不出来。def pde_residual(model, x): x x.clone().requires_grad_(True) u model(x) 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] residual -u_xx - (torch.pi ** 2) * torch.sin(torch.pi * x) return residual为什么不直接对原始 x 设requires_grad_因为训练循环里每次采样都会生成新的 x而 model 的输入要求对当前计算图可导。用 clone 再设置梯度标记可以避免把梯度标记残留到其他变量上也方便后续处理。4.3 训练循环与结果判断model PINN() optimizer optim.Adam(model.parameters(), lr1e-3) for epoch in range(3000): optimizer.zero_grad() x_coll torch.rand(128, 1) # 配置点随机采样 loss_pde torch.mean(pde_residual(model, x_coll) ** 2) x0 torch.zeros(32, 1) # 边界点多取几个更稳 x1 torch.ones(32, 1) loss_bc torch.mean(model(x0) ** 2) torch.mean(model(x1) ** 2) loss loss_pde loss_bc loss.backward() optimizer.step() if epoch % 500 0: print(epoch, loss.item())配置点每次迭代重新随机采样好处是网络接触的样本更丰富不容易过拟合到固定网格。边界条件用多组点而不是固定两个点也是为了让 BC 损失在优化过程中更稳定。训练完验证import numpy as np x_test torch.linspace(0, 1, 101).reshape(-1, 1) u_pred model(x_test).detach().numpy().ravel() u_true np.sin(np.pi * x_test.numpy().ravel()) print(最大误差:, np.max(np.abs(u_pred - u_true)))如果没有大问题这个算例最大误差通常能到 1e-3 量级。误差到 1e-2 也算基本学对。如果误差一直停在 0.1 以上说明训练没跑稳先看下面的排查章节。5. 同一道题用 TensorFlow 怎么写同一个方程TensorFlow 写法和 PyTorch 的差别很明显正好用来理解两个框架的自动微分机制。5.1 GradientTape 实现二阶导数import numpy as np import tensorflow as tf model tf.keras.Sequential([ tf.keras.layers.Dense(20, activationtanh), tf.keras.layers.Dense(20, activationtanh), tf.keras.layers.Dense(20, activationtanh), tf.keras.layers.Dense(1), ]) optimizer tf.keras.optimizers.Adam(learning_rate1e-3) def pde_residual(x): with tf.GradientTape(persistentTrue) as tape: tape.watch(x) u model(x) u_x tape.gradient(u, x) u_xx tape.gradient(u_x, x) residual -u_xx - (np.pi ** 2) * tf.sin(np.pi * x) return residualtape.watch(x)不能省。如果 x 是一个普通张量不 watch 的话 tape 不会追踪对它的求导tape.gradient(u, x)会返回 None。persistentTrue的作用前面说过是要支持多次调用 gradient。5.2 训练循环for epoch in range(3000): with tf.GradientTape() as tape: x_coll tf.random.uniform((128, 1)) loss_pde tf.reduce_mean(tf.square(pde_residual(x_coll))) x0 tf.zeros((32, 1)) x1 tf.ones((32, 1)) loss_bc (tf.reduce_mean(tf.square(model(x0))) tf.reduce_mean(tf.square(model(x1)))) loss loss_pde loss_bc grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) if epoch % 500 0: print(epoch, loss.numpy())这里最需要注意的是变量作用域tape.watch(x)写在计算导数的with块内部而优化参数梯度用另一层with tf.GradientTape()。如果两层 tape 混在一起容易把物理残差计算重复包进去造成不必要的开销甚至梯度错误。5.3 两版代码的核心差异对照步骤PyTorchTensorFlow网络定义nn.Moduletf.keras.Sequential输入跟踪requires_grad_(True)tape.watch(x)一阶导torch.autograd.gradtape.gradient二阶导再调用一次 gradpersistent tape 再调用一次参数梯度loss.backward()tape.gradient(loss, trainable_variables)参数更新optimizer.step()apply_gradients看明白这张表两套代码就可以互相翻译了。后面课程里遇到复杂算例无论原作者用哪个框架你都能换成自己熟悉的写法。6. 调参、判据和排查顺序PINN 训练不是只要跑了就能收敛。很多问题不是框架 bug而是参数设置或损失函数构成不合理。下面给出判断标准和排查顺序。6.1 判断训练是否正常的三个信号第一看总 loss 是否持续下降。如果前几百步 loss 不降反升优先调学习率或检查数据范围。第二把 loss 拆开看。不要只盯着总和PDE 残差和 BC 损失要分开打印。如果 PDE 降了但 BC 不降说明边界权重太低反过来说明边界权重压过了物理项。第三用已知解析解验证。上面的 sin(πx) 算例特别适合做基线因为误差能直接量化。6.2 常见收敛问题和解决顺序出现问题时按下面顺序排查先看输入坐标范围。输入 x 在 [0,1] 和 [0,1000] 是完全不同的训练难度。范围过大时网络很难学习先做归一化。再看导数是否为 None。PyTorch 里检查create_graphTrue是否开启TensorFlow 里检查tape.watch(x)是否调用。这一步几乎能解决一半的“二阶导报错”问题。接着看学习率。PINN 里 Adam 的默认 1e-3 通常没问题但 loss 震荡厉害时降到 1e-4 试一下。然后看网络宽度和深度。20 乘 3 的配置只适合简单方程。方程复杂或解的梯度很陡时要加深网络或换激活函数。最后看损失权重。BC 和 PDE 量级差太多时要么手动加权重要么把每一项做归一化。有时候 loss 一直降但测试误差很大这不是调参能解决的说明网络“记住”了配置点但没学会物理规律。此时增加配置点数量、增加采样密度或者改用更稳定的训练策略比如 L-BFGS 做最后阶段优化。6.3 批量算例的工程化思路课程后面的作业里很可能会连续跑多个方程或多种参数。这时候不要一个脚本改来改去建议把模型定义、残差函数、训练循环拆成函数用配置字典控制方程参数。config { equation: poisson, n_collocation: 128, n_boundary: 32, epochs: 3000, lr: 1e-3, hidden_layers: [20, 20, 20], }每次跑完把 loss 曲线、测试误差和模型参数保存到独立目录。文件名里带上算例标识和时间戳避免覆盖。日志里同时输出每个 loss 分量方便复现问题。凡是批量跑多个算例一定要把日志、输出目录和模型保存路径提前设计好。临时改文件名十次里有八次要出错。7. 框架之外PINN 入门真正要盯住的边界最后说几句容易被忽略的边界问题。PyTorch 和 TensorFlow 都能跑通 PINN但 PINN 本身有适用范围。简单方程、平滑解、规则边界基本 PINN 很容易收敛。碰到高频率振荡解、多尺度问题、强非线性或者参数突变默认配置的 PINN 大概率效果很差需要引入改进结构比如分区训练、特征缩放、自适应权重等。这些是后续课程的内容但你现在就要有这个预期否则会误以为框架选错了。另外如果你不想每次从零搭模型可以直接用 DeepXDE 这类封装库它在 PyTorch 和 TensorFlow 后端上都支持。我的建议是先用原始框架写一遍理解内部机制再考虑封装库。直接上库会少踩很多坑但也少学到积分点采样、残差构造和梯度回传这些核心细节。踩过几轮之后你会发现PINN 入门最耗时间的不是数学而是三件事环境没配对、自动微分用法不对、损失各项权重失衡。框架只是载体真正要掌握的是如何把物理残差变成可优化目标。先把 PyTorch 版的 sin(πx) 算例跑通再对照 TensorFlow 版写一遍这一讲的核心任务就算完成了。后续复杂算例其实都是在同一个代码骨架上加东西差别没有想象中大。
返回列表