PINN入门：PyTorch与TensorFlow框架选型及最小实现

物理信息神经网络PINN入门课程走到第 11 讲这一讲不写复杂算例先把框架选型这件决定后面所有代码风格的事情讲清楚。前面的课程已经解释了 PINN 的核心思路把偏微分方程PDE的残差、边界条件、初值条件一起塞进损失函数用神经网络反向传播去逼近物理系统的解。到了真正写代码这一步新手面临的第一道坎不是方程而是选 PyTorch 还是 TensorFlow。这个选择题没有标准答案但不同选择会直接影响你后续写损失函数、做自动微分、跑批量算例、甚至部署到工程环境的效率。这一讲会把两个框架在 PINN 场景下的真实差异拆开对比给出可复制的环境搭建命令分别用 PyTorch 和 TensorFlow 写一个最小可运行的 PINN 代码然后告诉你如何验证训练是否真正收敛、遇到 loss 不下降和梯度异常时该从哪里排查。读完这一讲你应该能根据自己手里的 GPU、后续应用场景和代码习惯做出一个不后悔的框架选择。1. 核心能力速览PINN 代码本质上只需要三个能力神经网络前向计算、自动微分求物理残差、优化器更新参数。两个主流框架都具备但使用方式和工程侧重点不同。选型维度PyTorch 方案TensorFlow 方案说明自动微分方式torch.autograd.grad动态计算图tf.GradientTape默认 Eager 执行PINN 需要高阶导数两者都能实现写法差异明显网络构建nn.Module子类化Python 风格keras层 API 或自定义tf.keras.ModelPyTorch 更贴近 Python 原生习惯TensorFlow 更偏工程化组装自定义损失直接写函数返回标量 tensor 即可继承tf.keras.Model重写train_step或用GradientTape手动更新PINN 的损失函数包含 PDE 残差两个框架都支持调试体验print直接打印中间张量断点调试友好2.x 后 Eager 模式也支持但部分 1.x 旧代码风格有遗留新手建议优先考虑调试便利性部署环境TorchScript / ONNX / TorchServeTensorFlow Serving / Lite / ONNX如果后续要把 PINN 模型做成接口服务需要提前考虑学习资源PINN 论文附带的官方代码很多基于 PyTorch老牌 PINN 库和部分工业案例使用 TensorFlow找到参考代码的速度决定上手效率从 PINN 入门角度说更稳妥的建议是如果你过去没写过任何深度学习代码优先选 PyTorch。它的动态图和 Python 原生风格能让你把注意力放在“物理损失怎么写”上而不是被框架语法绊住。如果你所在团队已经有 TensorFlow 部署链路或者后续要把 PINN 模型接入工业流水线就选 TensorFlow。2. PINN 为什么必须用框架先回答一个很多初学者会问的问题PINN 不就是解方程吗为什么不能自己写个普通全连接网络普通的神经网络训练损失函数一般只依赖模型输出和标签之间的差值比如交叉熵或均方误差。但 PINN 的损失函数里包含对模型输出的导数项。以最简单的二维热传导方程为例损失函数要计算温度场对空间坐标的二阶导数、对时间的导数然后把方程残差作为一项加到损失里。这个需求直接指向深度学习框架最核心的能力自动微分。你可以手动推导导数表达式再用有限差分近似但那样精度低、实现复杂而且边界条件一旦复杂就基本不可维护。自动微分能把“网络输出对输入求导”这件事在每次前向传播中自动完成无论是二阶导、混合偏导还是几十个输出分量一起求导都只需要几行代码。PINN 对框架的诉求可以总结为三点任意阶自动微分物理方程里的残差项通常是高阶导框架必须支持连续求导例如先求一阶导再对一阶导求导得到二阶导。灵活的自定义损失PINN 的损失由 PDE 残差、边界条件、初值条件、数据拟合项等组成每项权重不同框架不能限制你只能使用预置损失函数。GPU 加速真实物理问题往往需要成千上万次迭代纯 CPU 训练会很慢框架必须能很方便地把张量运算切到 CUDA。PyTorch 和 TensorFlow 都能满足这三个诉求差别在于写起来顺不顺手。接下来的章节分别演示。3. PyTorch 还是 TensorFlow按场景选3.1 什么时候选 PyTorch如果你是学术研究或课程学习场景选择 PyTorch 更合适。原因是近几年发布的 PINN 相关论文中基于 PyTorch 的参考代码占比很高。遇到问题你很容易在 GitHub 上找到对应实现把别人的网络结构、损失函数写法直接拿过来改。PyTorch 的调试体验对新手非常友好。你可以在任意一行代码处设置断点查看某个中间张量的形状和数值。PINN 训练过程中经常需要检查某个物理量残差的具体数值这个特点会大幅节省排查时间。另外PyTorch 的生态组件对自定义模型非常友好。torch.utils.data可以做数据加载和 batch 划分torch.optim提供 Adam、L-BFGS 等优化器后面如果要做批量物理参数扫描组合起来很方便。3.2 什么时候选 TensorFlow如果团队已经有 TensorFlow 部署基础设施或者模型最终要放到服务端长期运行TensorFlow 的部署链路会更成熟。TensorFlow Serving 对模型的版本管理、线上更新、并发请求处理都有完整方案这在工业场景中比“代码写起来是否顺手”更重要。另外TensorFlow 的tf.keras高层 API 封装程度很高网络层定义、训练循环结构比较固定。如果你只做标准的前馈网络拟合物理场用 Keras Sequential API 几行就能搭完一个网络代码量看起来确实更简洁。需要提醒的是TensorFlow 2.x 虽然默认 Eager 执行自定义 PINN 时会用到tf.GradientTape和手动更新参数这和 PyTorch 的autograd思考方式有些差异但学会之后也能顺畅使用。3.3 我的建议这一讲给出的结论是入门阶段优先选 PyTorch不要在这里过度纠结。核心原因是 PINN 学习过程本身就是不断修改损失函数和网络结构来对照物理现象的过程PyTorch 的动态图和 Python 风格能让你用最低的认知成本完成这些事情。等真正进入工程项目如果框架满足不了性能或部署需求再考虑迁移到 TensorFlow 或其他推理框架。学习阶段的代码迁移成本不高因为物理损失的计算逻辑是通用的换框架只需要重写网络和求导部分。4. 本地部署环境准备无论选哪个框架先准备好 Python 环境和 GPU 环境。下面给出通用步骤适用于 Windows 和 Linux具体版本号请以官方文档为准。4.1 创建独立虚拟环境PINN 学习中会频繁安装和更新依赖强烈建议使用虚拟环境隔离不要直接装到系统 Python 里。# 使用 conda 创建 Python 3.10 环境 conda create -n pinn python3.10 conda activate pinn如果不使用 conda用 venv 也可以python -m venv pinn_env source pinn_env/bin/activate4.2 安装 PyTorchPyTorch 的安装命令需要根据你的操作系统、CUDA 版本和是否使用 GPU 来选择建议直接访问 PyTorch 官网获取对应命令。一个通用模板如下# CPU 版本 pip install torch # GPU 版本示例实际安装命令以官网选择为准 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121注意在命令行直接运行pip install torch默认安装的是 PyTorch 官网的预编译包CPU 版和 GPU 版的差别在安装后可以通过torch.cuda.is_available()验证。import torch print(torch.__version__) print(torch.cuda.is_available())如果输出False说明当前环境没有可用的 CUDA训练时只能走 CPU 路径。4.3 安装 TensorFlowTensorFlow 同样分为 CPU 版和 GPU 版。安装前确认你的显卡驱动支持 CUDA 版本一个通用模板如下# CPU 版本 pip install tensorflow # GPU 版本通过官方 pip 源安装 pip install tensorflow安装完成后同样需要验证 GPU 是否可用import tensorflow as tf print(tf.__version__) print(tf.config.list_physical_devices(GPU))如果 GPU 列表为空说明 TensorFlow 没有识别到 GPU需要检查驱动和 CUDA 工具包版本。4.4 硬件与磁盘PINN 训练对 GPU 显存的要求与网络规模和批大小有关。入门阶段用全连接网络、每层 20 到 64 个神经元、批大小 128 左右时显存占用并不高很多集成显卡或低端独立显卡也能运行。但要注意如果做高维 PDE 或大批量训练显存占用会快速上升。这不是框架本身的问题而是自动微分需要保存反向传播所需的中间张量。降低显存占用的通用办法是减小 batch size、减少网络层数、降低输入分辨率或使用混合精度训练。磁盘空间方面Python 环境加两个框架包大约需要 3 到 5 GB如果包含 CUDA 工具包会更多。建议为训练实验单独准备一个数据目录把输入数据、模型权重和输出结果分目录管理。5. 用 PyTorch 写一个最小 PINN下面用一维泊松方程作为示例展示 PyTorch 中 PINN 的完整代码结构。方程形式为[ -u(x) \pi^2 \sin(\pi x), \quad x \in [0, 1] ]边界条件为 ( u(0)0, u(1)0 )精确解为 ( u(x)\sin(\pi x) )。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, 1) ) def forward(self, x): return self.net(x) # 初始化模型和优化器 model PINN() optimizer torch.optim.Adam(model.parameters(), lr1e-3) # 计算方程残差损失 def pde_loss(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] # 方程右端项 f torch.pi**2 * torch.sin(torch.pi * x) # 方程残差 residual -u_xx - f return torch.mean(residual**2) # 边界条件损失 def boundary_loss(): x_left torch.zeros((100, 1), requires_gradFalse) x_right torch.ones((100, 1), requires_gradFalse) u_left model(x_left) u_right model(x_right) return torch.mean(u_left**2) torch.mean(u_right**2) # 训练循环 n_iter 5000 for it in range(n_iter): optimizer.zero_grad() # 内部配置点用于计算方程残差 x_collocation torch.rand((256, 1)) # 在 [0,1] 内采样 loss_pde pde_loss(x_collocation) loss_bc boundary_loss() # 加权组合 loss loss_pde loss_bc loss.backward() optimizer.step() if it % 500 0: print(fIter {it}, Loss: {loss.item():.6f})这段代码的关键点有三个requires_grad_(True)让输入张量参与自动微分这样神经网络输出才能对输入求梯度。torch.autograd.grad中的create_graphTrue表示允许对梯度再求梯度这是计算二阶导的必要条件。边界条件通过直接输入x0和x1的采样点来约束模型输出。只有一阶导的方程create_graph可以设置为 False需要二阶或更高阶导数时必须保留计算图。6. 用 TensorFlow 写一个最小 PINNTensorFlow 版本使用GradientTape完成相同的计算。import tensorflow as tf # 定义网络 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(1) def call(self, x): x self.dense1(x) x self.dense2(x) return self.dense3(x) model PINN() optimizer tf.keras.optimizers.Adam(learning_rate1e-3) # 计算方程残差损失 def pde_loss(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) f tf.constant(tf.math.pi**2, dtypetf.float32) * tf.sin(tf.constant(tf.math.pi, dtypetf.float32) * x) residual -u_xx - f return tf.reduce_mean(tf.square(residual)) # 训练循环 for it in range(5000): with tf.GradientTape() as tape: x_collocation tf.random.uniform((256, 1), minval0.0, maxval1.0) # 方程残差损失 loss_pde pde_loss(x_collocation) # 边界条件损失 x_left tf.zeros((100, 1)) x_right tf.ones((100, 1)) loss_bc tf.reduce_mean(tf.square(model(x_left))) tf.reduce_mean(tf.square(model(x_right))) loss loss_pde loss_bc grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) if it % 500 0: print(fIter {it}, Loss: {loss.numpy():.6f})TensorFlow 版本的差异点tf.GradientTape(persistentTrue)允许重复调用多次梯度操作。求完一阶导再求二阶导时如果用普通GradientTape需要在同一个上下文内完成全部计算。tf.random.uniform生成采样点张量默认参与自动微分但通过tape.watch(x)显式监听输入张量。参数更新需要手动apply_gradients这一步在 PyTorch 里对应optimizer.step()。从代码量看两个框架差别不大。真正影响体验的是你对哪种写法的直觉更接近。PyTorch 的tensor.grad和autograd更透明TensorFlow 的GradientTape更强调作用域。7. PINN 典型问题测试与结果验证写完最小代码后下一步是验证训练是否真的收敛到了物理上正确的解。7.1 用精确解对照上面的泊松方程有解析解 ( u(x)\sin(\pi x) )。训练结束后在 [0,1] 上均匀取 100 个点对比网络输出和解析解import numpy as np import torch # PyTorch 对比示例 model.eval() x_test torch.linspace(0, 1, 100).reshape(-1, 1) u_pred model(x_test).detach().numpy().ravel() u_true np.sin(np.pi * x_test.numpy().ravel()) # 计算最大误差和均方误差 error np.abs(u_pred - u_true) print(fMax error: {error.max():.2e}) print(fRMSE: {np.sqrt(np.mean(error**2)):.2e})一般来说迭代到后期最大误差降到 ( 10^{-2} ) 到 ( 10^{-3} ) 量级说明模型已经学习到了方程的主要特征。如果一直停留在 ( 10^{-1} ) 以上需要检查网络层数、激活函数、采样点数和迭代次数。7.2 观察损失曲线的不同阶段PINN 训练中损失曲线通常不会像图像分类那样平滑下降。物理残差和边界条件损失在初始阶段权重不同可能出现以下几个现象边界条件损失快速下降但 PDE 残差损失下降缓慢。这说明网络优先拟合了边界内部物理约束还没学透。可以适当增大迭代次数或调整两项损失的权重。PDE 残差下降到某个值后不再变化。此时不一定是代码写错了而是网络容量或采样点密度不足。尝试加深网络、增加采样点或者切换到 L-BFGS 这类二阶优化器。训练后期损失出现周期性波动。这是配置点每次重新采样导致的属正常现象。更稳定的做法是把配置点固定并进行较多次迭代或者使用带重置策略的采样方法。7.3 判断收敛的物理准则PINN 训练不仅看 Loss 数值更要看预测解是否满足物理规律。以热传导方程为例预测的温度场不应该出现局部突变或违背能量守恒。因此测试时应增加一组物理约束验证检查预测值的单调性、极值位置、边界值误差。如果预测结果在某些区域明显偏离物理直觉优先怀疑配置点分布是否覆盖了这些区域。PINN 以配置点方式离散物理域配置点稀疏的地方网络学习到的物理约束就弱。8. 资源占用与性能观察PINN 训练的资源占用没有固定数字因为它与网络大小、批量大小、方程维度和自动微分阶数直接相关。不过可以给出通用的观察和调优思路。8.1 如何观察显存和内存占用训练过程中可以用系统自带的nvidia-smi命令查看 GPU 显存占用或者使用watch -n 1 nvidia-smi持续刷新。在 Python 内部也可以用torch.cuda.memory_allocated()获取当前显存占用。print(torch.cuda.memory_allocated() / 1024**2, MB)如果显存不足优先减少 batch size。PINN 使用随机采样batch size 从 256 降到 128 几乎不影响物理规律学习反而可能让单次迭代的稳定性提升。8.2 CPU 和 GPU 训练差异PINN 入门代码在没有 GPU 的环境中也能跑起来只是训练速度慢。CPU 训练 5000 次迭代可能在几分钟到十几分钟GPU 训练通常能在几秒到几十秒内完成。差异主要来自矩阵运算的并行度和自动微分构建计算图的开销。如果你暂时没有 GPU优先用小网络小采样点跑通流程。但要注意最后的效果验证部分必须基于真实训练结果不能只用小采样点数量“硬推”出物理规律。8.3 如何降低显存占用几个常见手段减少网络层数和每层神经元数。降低每个 iteration 的采样点数量。使用混合精度训练例如 PyTorch 的torch.autocast。清理不再使用的中间张量例如del u_x, u_xx; torch.cuda.empty_cache()。这些手段在入门阶段不一定都需要但它们能帮助你理解 PINN 的资源瓶颈在哪里后续做高维问题时直接用得上。9. 常见问题与排查方法PINN 训练中的问题很多下面列出最常遇到的几类。问题现象可能原因排查方式解决方案安装框架后import torch报错环境类型不匹配、Python 版本不兼容检查 Python 版本和 pip 安装源重新创建虚拟环境按官方安装命令安装torch.cuda.is_available()返回 FalseCUDA 驱动版本过低或 PyTorch 装成 CPU 版运行nvidia-smi查看驱动版本安装与驱动匹配的 CUDA 版 PyTorchTensorFlow 识别不到 GPUCUDA、cuDNN 版本与 TensorFlow 版本不匹配运行tf.config.list_physical_devices(GPU)参照官方 GPU 支持矩阵安装对应版本训练时 loss 不下降学习率过大或过小、网络结构不合理、损失函数权重失衡打印各分项损失数值调小学习率、调整损失权重、加大采样点数二阶导计算结果异常create_graph参数错误或采样点超出定义域打印中间梯度值比较检查create_graphTrue是否设置训练后期 loss 波动大配置点随机采样导致梯度噪声观察波动幅度固定配置点或增大 batch size预测边界值偏差大边界条件权重太低检查边界损失数值提高边界条件损失权重模型在局部区域预测异常配置点稀疏可视化采样点分布该区域加密采样点针对“训练时 loss 不下降”这个最常见问题一个有效的排查思路是先把物理损失去掉只保留边界条件损失看模型是否能学会边界值。如果能学会再逐步加回物理残差定位是哪一项导致数值不稳定。10. 框架学习路径与工程实践建议选框架只是起点真正决定 PINN 学习进度的是后面几个习惯。第一建议把“最小可运行代码”保存下来。后续所有更复杂的算例都在这份代码基础上修改网络、损失函数和数据加载逻辑。不要每次从零开始写。第二训练过程中要有记录意识。每次修改网络结构、采样点数或学习率后记录 loss 曲线和最终预测误差并保留模型权重文件。这样能准确判断哪个改动带来了实际效果而不是靠感觉调参。第三批量实验时要做好任务队列。PINN 经常需要扫描多个物理参数例如扩散系数、边界温度、材料属性等。可以先写一个参数列表循环创建不同实验目录每个目录下保存对应配置和模型输出。接口调用层面如果后面要把训练好的 PINN 模型封装成服务优先考虑 ONNX 导出或框架自带的 serving 方案这样其他程序可以通过统一接口访问模型。第四合规问题需要从一开始就重视。如果使用实验数据、仿真数据或任何第三方数据集务必确认数据来源合法、具备使用权。涉及真实工业生产数据、人体相关数据或受版权保护的材料时必须做匿名化和授权合规处理。PINN 的应用场景越多数据合规的边界就越重要。第五不要盲目使用最新版本框架。对于 PINN 学习稳定性和生态兼容性比新功能更重要。热门框架版本升级后自动微分 API 或部署工具链可能发生变动优先选择文档较为成熟、周围同学或同事用得多的版本。11. 这一讲之后第 12 讲方向本讲已经完成框架选型和最小代码验证。接下来第 12 讲会继续深入 PINN 的损失函数权重配置、激活函数选择以及边界条件的标准化处理这些都是实际训练中直接影响收敛速度和精度的细节。建议你先动手跑一遍本讲的代码把 PyTorch 和 TensorFlow 两个版本都运行一次对比两者的 loss 下降速度。通过这个对比你才能真正理解框架选择对你编程习惯和调试效率的影响。选择 PyTorch 还是 TensorFlow答案不在别人嘴里在你自己的运行结果里。