PyTorch与TensorFlow实现PINN：一维泊松方程实操对比

PINNPhysics-Informed Neural Networks物理信息神经网络在最近几年的工程应用里越来越常见很多做流体、固体、热传导、电磁场甚至飞行轨迹预测的同学都在尝试把物理方程直接放进神经网络的损失函数里。但真正动手时大多数人遇到的第一个问题不是方程怎么离散而是框架选 PyTorch 还是 TensorFlow以及框架装好之后如何用自动微分把二阶导数写对。作为 PINN 入门系列课程的第 11 讲这篇不堆理论而是从环境配置开始把一个一维泊松方程的最小 PINN 案例分别用 PyTorch 和 TensorFlow 跑通并把两者的网络定义、求导方式、训练循环、常见报错和工程选型放在一起对比。读完应该能对“PINN 到底要框架做什么”有一个完整认知也能直接拿这个最小案例去扩展自己的方程。1. 先理解 PINN 的工作机制再谈框架选型1.1 PINN 在做什么让神经网络去拟合物理方程普通神经网络训练时损失函数通常来自数据标签比如分类用交叉熵、回归用均方误差。PINN 的差别在于除了数据驱动它还把物理方程本身当作约束。假设要求解一个微分方程F(u, ∂u/∂x, ∂²u/∂x²) 0, x ∈ Ω边界条件为 B(u) 0, x ∈ ∂Ω。PINN 的做法是构造一个神经网络 u_θ(x)其中 θ 是待学习的参数然后定义损失函数L L_pde λ_bc · L_bc其中 L_pde 是方程残差的均方误差也就是把网络输出代入方程后左边不等于 0 的程度L_bc 是边界条件的误差。训练完成后网络输出 u_θ(x) 就近似满足方程和边界条件。这里最关键的一步是计算网络输出对输入 x 的导数。二阶甚至更高阶导数在方程里很常见所以框架的自动微分Automatic DifferentiationAD能力直接决定了 PINN 代码的写法。可以说PINN 的代码复杂度大部分不在网络结构而在“如何把导数写对”。1.2 框架在 PINN 中承担哪几层职责把 PINN 训练过程拆开看框架至少承担五层职责网络结构MLP、卷积、残差连接等层定义和参数管理。自动微分计算 u 对 x 的一阶、二阶导数这是 PINN 区别于普通监督学习的核心。优化器Adam、L-BFGS 等参数更新算法。设备调度CPU/GPU 的张量计算和显存管理。模型保存、加载、推理和部署。PyTorch 和 TensorFlow 在这五层都能完成但写法和生态侧重点不同。选型时不是看“哪个更好”而是看“哪个更适合你的方程、团队和部署环境”。1.3 PyTorch 和 TensorFlow 对 PINN 的关键能力差异PyTorch 的自动微分基于动态计算图每次前向传播都会重新建立计算图使用torch.autograd.grad可以非常直观地取出任意中间变量对输入的导数。TensorFlow 在 eager 模式下使用tf.GradientTape记录前向过程再调用gradient得到导数。求二阶导数时PyTorch 需要create_graphTrueTensorFlow 需要嵌套两层GradientTape。两者难度都不高但初次接触时容易在 API 细节上卡住。能力维度PyTorchTensorFlow自动微分入口torch.autograd.gradtf.GradientTape计算图模式动态图为主eager 模式为主也可静态图二阶导数写法create_graphTrue两层GradientTape嵌套自定义训练循环通常手写for循环直观tf.GradientTape内手写也可用 Keras社区与论文复现科研论文默认选择多工业界存量项目多部署生态TorchScript、TorchServeTF Serving、TFLite链路更完整与 PINN 库结合DeepXDE 可选 PyTorch 后端DeepXDE 原生支持 TF 后端从 PINN 角度看两种框架都能完成同样的事情真正影响体验的是求导 API 的书写习惯以及后续要接入的库是哪个生态。2. 环境准备PyTorch 与 TensorFlow 安装注意事项很多 PINN 新手不是死在方程推导而是死在环境安装。PyTorch 和 TensorFlow 对 CUDA 版本、Python 版本、驱动版本都有要求装错之后常见的现象是import torch成功但torch.cuda.is_available()返回False或者 TensorFlow 能装上但训练时根本没用上 GPU。2.1 先确认 Python、CUDA、GPU 驱动版本安装之前先用下面三条命令确认本机状态python --version nvidia-smi nvcc --version这里要澄清一个常见误解nvidia-smi显示的是当前 GPU 驱动支持的最高 CUDA 版本并不代表你已经安装了对应版本的 CUDA 工具包nvcc --version才是已安装 CUDA 工具包的版本。PyTorch 和 TensorFlow 的 pip 包通常自带运行时所需的 CUDA 库不一定需要单独安装完整 CUDA 工具包但驱动版本必须足够新。如果本机没有 NVIDIA GPU或者只是在学习阶段先用 CPU 版本跑通流程也完全可行。PINN 的最小案例计算量不大CPU 上几分钟就能收敛。生产环境或大规模方程再考虑 GPU。2.2 使用虚拟环境安装 PyTorch推荐用 Anaconda 创建独立虚拟环境避免把系统 Python 环境搞乱conda create -n pinn python3.10 -y conda activate pinnCPU 版本直接安装pip install torch torchvision torchaudioGPU 版本需要到 PyTorch 官网选择 CUDA 版本然后执行对应命令。以 CUDA 12.1 为例安装命令形如pip install torch --index-url https://download.pytorch.org/whl/cu121这里要特别说明上面的cu121只是示例实际安装前要到官网安装页面确认你需要的 CUDA 版本和 Python 版本。如果使用昇腾 NPU 等特定硬件平台还需要到对应厂商提供的安装源获取适配版本。不要照抄网上任意一条命令版本不匹配是 PINN 环境问题的主要来源。torchvision和torchaudio对 PINN 本身不是必需的但保持同一套版本安装可以避免后续引入其他视觉任务时产生依赖冲突。2.3 使用虚拟环境安装 TensorFlowTensorFlow 的安装相对简单但版本策略变动较多落地前一定要看官方文档。常见的安装命令是pip install tensorflow需要 CUDA 依赖时TensorFlow 官方也提供了一键安装方式pip install tensorflow[and-cuda]需要注意TensorFlow 在 Windows 原生环境的 GPU 支持策略在不同版本之间有差异官方推荐在 Linux 或 WSL2 下使用 GPU 版本。如果只是在 Windows 上学习验证建议先跑 CPU 版本把 PINN 的求导和训练逻辑确认好再迁移到 Linux 环境做大规模计算。2.4 安装完成后必须做一次 GPU 可用性检查环境装完不要直接写 PINN 代码先检查框架是否真的能调用 GPUpython -c import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))python -c import tensorflow as tf; print(tf.__version__, tf.config.list_physical_devices(GPU))正常的输出应该看到版本号、True以及 GPU 设备名。如果输出False或者空列表通常说明 CUDA 驱动版本过低、框架版本不匹配或者安装的是 CPU 版。注意框架能import成功不代表能用 GPU。torch.cuda.is_available()和tf.config.list_physical_devices(GPU)才是判断依据。3. 用 PyTorch 实现最小 PINN一维泊松方程环境就绪后先用一个最简单但有完整物理意义的方程跑通流程。这里选择一维泊松方程因为它有精确解方便验证 PINN 收敛质量。3.1 问题定义和精确解求解区间为 x ∈ (0, 1)方程如下u(x) -π² · sin(πx)边界条件为 u(0) 0u(1) 0。这个方程的精确解是 u(x) sin(πx)所以训练完成后可以直接对比网络输出和精确解计算最大误差和均方误差。PINN 的损失函数由两部分组成PDE 残差损失将网络输出代入方程计算 u_θ(x) π²·sin(πx) 的均方误差。边界条件损失计算 u_θ(0) 和 u_θ(1) 的均方误差。3.2 网络结构与激活函数选择PINN 最常见的网络结构是 MLP隐藏层用 Tanh 激活。这里有一个关键点微分方程需要二阶导数如果隐藏层使用 ReLU二阶导数在绝大多数区域恒为 0网络就没有能力逼近 u(x) 的非零值训练必然失败。Tanh、Sigmoid、Swish 这类光滑激活函数二阶导数存在且非平凡更适合 PINN。网络结构写成import torch import torch.nn as nn 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)输入是 x输出是 u_θ(x)。这个结构没有任何技巧但对一维问题已经足够。3.3 二阶导数的自动微分写法PyTorch 求二阶导数核心是torch.autograd.grad的create_graphTrue参数。第一步求一阶导数时如果不设置create_graphTrue计算图会被释放后续无法再求二阶导数。def pde_loss(model, x): x.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 torch.mean(residual ** 2)这里grad_outputs是链式法则里的上游梯度因为 u 和 x 都是形状为 (N, 1) 的张量所以用torch.ones_like(u)。[0]是因为grad返回的是一个元组。3.4 完整训练代码边界条件损失、采样点更新和训练循环如下import numpy as np model PINN() optimizer torch.optim.Adam(model.parameters(), lr1e-3) def bc_loss(model, x_bc): u_bc model(x_bc) return torch.mean(u_bc ** 2) x_bc torch.tensor([[0.0], [1.0]]) for epoch in range(3000): # 每个 epoch 重新随机采样避免过拟合到固定网格 x_collocation torch.rand(256, 1) optimizer.zero_grad() loss_pde pde_loss(model, x_collocation) loss_bc bc_loss(model, x_bc) loss loss_pde loss_bc loss.backward() optimizer.step() if epoch % 500 0: print(fepoch {epoch}, loss_pde{loss_pde.item():.2e}, floss_bc{loss_bc.item():.2e})训练结束后用下面的代码验证预测精度import matplotlib.pyplot as plt x_test torch.linspace(0, 1, 100).reshape(-1, 1) u_pred model(x_test).detach().numpy() u_exact np.sin(np.pi * x_test.numpy()) print(max error:, np.max(np.abs(u_pred - u_exact))) plt.plot(x_test.numpy(), u_exact, r-, labelexact) plt.plot(x_test.numpy(), u_pred, b--, labelPINN) plt.legend() plt.show()如果训练正常max error应该能达到 1e-2 到 1e-3 量级曲线基本重合。注意采样点放在循环内部每个 epoch 重新随机生成这是 PINN 的常见做法。如果采样点固定不变网络容易在有限网格点上过拟合中间区域的方程残差可能仍然很大。4. 用 TensorFlow 实现同一个 PINN同一个方程用 TensorFlow 再写一遍。目的是对比两个框架在求导和训练循环上的差异而不是判断谁优谁劣。4.1 tf.GradientTape 求二阶导的方式TensorFlow 中求二阶导数需要使用两层GradientTape。外层 tape 需要设置persistentTrue因为后面还要通过它计算二阶导数内层 tape 负责求一阶导数。import tensorflow as tf import numpy as np class PINN(tf.keras.Model): def __init__(self): super().__init__() self.dense1 tf.keras.layers.Dense(20, activationtanh) self.dense2 tf.keras.layers.Dense(20, activationtanh) self.dense3 tf.keras.layers.Dense(20, activationtanh) self.dense4 tf.keras.layers.Dense(1) def call(self, x): x self.dense1(x) x self.dense2(x) x self.dense3(x) return self.dense4(x)4.2 完整训练代码def pde_loss(model, x): with tf.GradientTape(persistentTrue) as tape_xx: tape_xx.watch(x) with tf.GradientTape() as tape_x: tape_x.watch(x) u model(x) u_x tape_x.gradient(u, x) u_xx tape_xx.gradient(u_x, x) del tape_xx residual u_xx (np.pi ** 2) * tf.sin(np.pi * x) return tf.reduce_mean(tf.square(residual)) model PINN() optimizer tf.keras.optimizers.Adam(learning_rate1e-3) for epoch in range(3000): x_collocation tf.random.uniform((256, 1)) x_bc tf.constant([[0.0], [1.0]]) with tf.GradientTape() as tape: loss_pde pde_loss(model, x_collocation) u_bc model(x_bc) loss_bc tf.reduce_mean(tf.square(u_bc)) 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(fepoch {epoch}, loss_pde{loss_pde.numpy():.2e}, floss_bc{loss_bc.numpy():.2e})这里要注意几点。第一外层GradientTape必须persistentTrue否则tape_xx.gradient调用一次后资源就被释放。第二用完要del tape_xx手动释放资源避免长时间训练时内存累积。第三tf.random.uniform在循环内生成新的采样点和 PyTorch 版本的行为一致。4.3 两种框架代码结构对比对比项PyTorchTensorFlow网络定义nn.Module子类tf.keras.Model子类一阶导数torch.autograd.grad(u, x, create_graphTrue)内层tape.gradient(u, x)二阶导数对一阶结果再调grad外层tape.gradient(u_x, x)反向传播loss.backward()optimizer.step()tape.gradient(loss, trainable_variables)apply_gradients采样点生成torch.randtf.random.uniform从代码量看两者几乎没有差别。真正的差别在学习曲线PyTorch 的loss.backward()隐藏了梯度流向初学者更容易理解TensorFlow 需要手动区分tape求的是网络参数的梯度还是网络输出对输入的梯度这两类梯度在 PINN 里同时存在容易混淆。5. 运行验证从损失曲线和预测结果判断模型是否收敛5.1 期望的输出和判断标准训练 3000 个 epoch 后正常情况下的输出大概呈现这样的趋势epoch 0, loss_pde8.31e00, loss_bc5.62e-01 epoch 500, loss_pde1.20e-02, loss_bc1.90e-03 epoch 1000, loss_pde3.50e-03, loss_bc2.80e-04 epoch 1500, loss_pde1.80e-03, loss_bc5.10e-05 epoch 2000, loss_pde1.10e-03, loss_bc2.30e-05 epoch 2500, loss_pde7.60e-04, loss_bc1.10e-05 epoch 3000, loss_pde5.40e-04, loss_bc8.00e-06判断标准不要只看总损失要分别看loss_pde和loss_bc。边界条件损失应该快速下降PDE 残差损失下降会慢一些这是正常现象。最终在测试点上的最大绝对误差在 1e-2 到 1e-3 量级说明模型已经学到方程的解。5.2 训练过程常见现象解读如果训练曲线不符合预期可以从以下几类现象判断问题方向现象可能原因处理方向loss_pde一开始就很大且不降采样点包含边界附近非光滑区域或学习率过大降低学习率到 1e-4或先固定采样点观察loss_bc很快到 0但loss_pde几乎不动边界条件主导了梯度方程残差被忽略给loss_pde乘以较大权重例如 10 或 100loss_pde震荡剧烈采样点每个 epoch 随机变化损失本身有噪声增大批大小或改用固定网格加随机扰动总损失很小但预测误差很大PINN 出现平凡解或近似零解检查激活函数是否光滑边界条件权重是否过低5.3 学习环境与生产环境的差异学习阶段CPU 跑这个最小案例完全够用重点是理解求导和训练循环。但进入生产环境后至少要补齐以下内容GPU 训练和混合精度方程规模大时显存和速度差异明显。损失权重自适应调整简单相加在复杂方程上很难收敛。模型保存、加载和版本管理训练好的 PINN 需要固化下来做推理。日志、监控和结果校验不只是打印 loss还要定期对比已知工况的精确解或数值解。回滚方案PINN 训练不稳定时能回到上一次可用模型。6. PINN 实现中的 6 个常见坑6.1 激活函数选择错误导致二阶导数恒为零现象训练正常执行但loss_pde下降缓慢甚至不降最终预测结果是一条直线或接近零的曲线。原因隐藏层使用 ReLUReLU 的二阶导数在正区间为 0负区间为 0仅在 x0 处未定义。对二阶微分方程网络根本无法表达非零的 u。解决换成 Tanh、Sigmoid、Swish 等光滑激活函数。PINN 的激活函数选型要优先考虑“导数是否非平凡”而不是“效果好不好”。6.2 采样点固定导致过拟合到网格点现象训练时 loss 很低但在非采样点的测试位置误差明显偏大。原因采样点集合在整个训练过程中固定不变网络只需在有限点上满足方程残差中间区域没有被约束。解决每个 epoch 重新随机采样或者采用固定网格加随机扰动的方式。这也是 DeepXDE 等库默认的做法。6.3 PDE 残差和边界条件损失量级失衡现象loss_bc降到 1e-5loss_pde还在 1e-1总损失被边界条件主导方程约束形同虚设。原因两类损失的量级天然不同直接相加时梯度被大项主导。解决给损失加权L λ_pde · L_pde λ_bc · L_bc。简单做法是固定权重例如 λ_pde10、λ_bc1复杂做法是训练过程中根据梯度统计自适应调整。先做固定权重再考虑自适应方案。6.4 PyTorch 没有设置 create_graph 导致二阶导数报错现象运行到u_xx求导时报错信息包含element 0 of tensors does not require grad或类似提示。原因第一次调用torch.autograd.grad时没有设置create_graphTrue计算图被释放后续无法继续求导。解决一阶求导时必须写create_graphTrue。TensorFlow 版本则是外层GradientTape必须设persistentTrue使用后手动del释放。6.5 PyTorch 2.6 之后 torch.load 的 weights_only 报错现象加载模型权重时出现类似Weights only load failed ...的报错或者提示weights_only参数相关的 warning。原因PyTorch 2.6 开始torch.load默认把weights_only设为True这是为了安全考虑防止加载 pickle 文件时执行任意代码。直接用旧方式加载包含非张量对象的 checkpoint 就会失败。解决保存模型时优先使用torch.save(model.state_dict(), path)加载时用model.load_state_dict(torch.load(path))。如果确实需要加载完整 checkpoint可以显式设置torch.load(path, weights_onlyFalse)但要确认文件来源可信。6.6 输入特征未归一化导致训练发散现象输入 x 的取值范围很大例如模拟区域是毫米到米级别训练过程中 loss 出现nan。原因PINN 的输入和输出尺度差异过大时网络权重初始化与梯度更新不匹配容易发散。这一点在物理量纲差异大的问题上尤其明显。解决把空间坐标归一化到 [0, 1] 或 [-1, 1]输出量大时也做量纲归一化。训练开始前先打印输入、输出和 loss确认量级合理再训练。7. 框架选型建议与 PINN 工程化清单7.1 什么时候选 PyTorch什么时候选 TensorFlow从当前工程实践看可以按以下逻辑选择如果你是科研方向、需要大量参考论文代码选 PyTorch。PINN 相关论文的官方实现使用 PyTorch 的比例更高复现和研究阻力小。如果你在工业系统里做模型部署并且团队已经熟悉 TF Serving 或 TFLite选 TensorFlow。TensorFlow 的部署链路更完整但前提是你的方程场景能容忍它的自定义训练循环复杂度。如果你不想自己写网络和求导直接用 DeepXDE 这类 PINN 框架。DeepXDE 支持多种后端上层 API 统一底层后端可以切换。如果只是学习 PINN 概念两个都装一遍把同一个方程各跑一次。这个对比过程对理解自动微分非常有帮助。选型没有绝对答案关键是团队维护能力和部署环境。不要让框架选择成为 PINN 落地的阻碍。7.2 PINN 从原型到生产的工程清单写一个可复用的清单每次开始新的 PINN 任务前逐项确认检查项确认内容方程形式明确最高阶导数、边界/初始条件、定义域范围激活函数是否为光滑激活函数二阶导数是否非平凡输入归一化空间坐标是否归一化物理量纲是否统一采样策略是否每个 epoch 重新采样边界点数量是否足够损失权重PDE 残差和边界条件是否分开展示权重是否可调导数写法PyTorch 是否create_graphTrueTF 是否persistentTrueGPU 验证cuda.is_available()或list_physical_devices是否为真结果校验是否有精确解、数值解或实验数据用于对比模型保存是否保存state_dict或完整模型是否记录训练参数和损失曲线日志监控是否记录每个 epoch 的分项损失是否能在训练异常时快速定位7.3 后续学习路径跑通一维泊松方程后建议按这个顺序扩展把一维方程换成二维 Laplace 方程或 Poisson 方程理解多变量求导。加入时间维尝试 Burgers 方程理解初始条件和时空采样。引入更复杂的方程比如 Navier-Stokes 方程或弹道轨迹参数辨识类问题理解多损失项平衡。尝试 DeepXDE 或 NVIDIA Modulus 这类专业 PINN 库把网络结构、采样、损失加权交给成熟工具自己专注方程和物理约束设计。学习领域自适应和权重点火这类进阶技巧它们能明显提升复杂方程的收敛质量。这一讲的核心结论是PINN 与框架的关系不在于框架本身多强大而在于自动微分 API 用得是否顺手。把 PyTorch 和 TensorFlow 的最小案例各写一遍二阶导数、采样点、损失权重的坑都踩过一遍之后再回到自己的方程思路会清楚很多。