TensorFlow中PINN迭代优化时创建可追踪梯度的新数组
2026-7-6
问题概述
在TensorFlow中构建物理信息神经网络（PINN），基础PINN训练正常，现需新增迭代优化模块：遍历预测结果的patch来细化预测值，需满足：
- 创建与原预测数组形状相同的新数组存储细化结果
- 不能覆盖原预测值（需用原预测值构建每个patch）
- 新数组需支持
tf.GradientTape的自动微分梯度追踪
已尝试方法及问题
- 每次迭代创建新
tf.Variable：触发ValueError，tf.function仅支持单例tf.Variables - 全局创建
tf.Variable并在每个epoch重置：模型优化时损失无下降，推测自动微分无法追踪操作（梯度为None） - 使用
tf.identity创建深拷贝：操作无法被梯度tape完整记录，模型无法优化 - 考虑
tf.TensorArray：不支持元素重赋值，且无法保证梯度追踪完整性
相关代码片段
基础组件与训练循环
def init_model(layers): '''初始化前馈神经网络的可训练部分：权重、偏置、激活函数。''' model = ... # 原初始化逻辑 return model def MSE(truths, preds): '''计算两个输入向量的均方误差，输出单个数值''' return tf.reduce_mean(tf.square(truths - preds)) def find_uv(psi, tape2, x, y): '''从流函数（神经网络的两个输出）中求解速度''' u, v = ... # 原求解逻辑 return u, v def find_dvtvs(p, u, v, tape2, x, y, t): '''求解控制方程残差所需的高阶导数''' u_x,v_x,p_x, u_y,v_y,p_y, u_t,v_t, u_xx,v_xx, u_yy,v_yy = ... # 原求解逻辑 return u_x,v_x,p_x, u_y,v_y,p_y, u_t,v_t, u_xx,v_xx, u_yy,v_yy def get_residual(p, u, v, u_x,v_x,p_x, u_y,v_y,p_y, u_t,v_t, u_xx,v_xx, u_yy,v_yy): '''计算二维Navier-Stokes方程的残差''' residual = ... # 原残差计算逻辑 return residual def train_step(): '''计算当前epoch的损失，以及损失相对于可训练参数的梯度''' loss, grad_theta, dataerror, physicserror = get_grad(model, dxy_coords, dxy_truths, hx_coords, hy_coords ) # 执行梯度下降步骤 optim.apply_gradients([ (grad, var) for (grad, var) in zip(grad_theta, model.trainable_variables) if grad is not None ]) return loss, dataerror, physicserror # 训练循环 layers = [3,10,10,2] #输入层为xyt，两个隐藏层各10个神经元，输出层为psi（流函数）和p（压力） N = 1000 #训练轮数 model = init_model(layers) optim = tf.keras.optimizers.Adam() # 假设优化器为Adam for i in range(N+1): loss, dataerror, physicserror = train_step()
核心问题所在的HCP函数
def HCP(puv, dxyshape, hxshape, hyshape ): '''将预测结果转换为约束patch，处理后得到单个值并输出处理后的结果''' # 提取各类数据的数量 Ndxy = dxyshape[0]*dxyshape[1]*dxyshape[2] #=500 Nhx = hxshape[0]*hxshape[1]*hxshape[2] #=400 Nhy = hyshape[0]*hyshape[1]*hyshape[2] #=400 # 提取p,u,v并重塑为3D数组 p_dxy = tf.reshape(puv[:Ndxy,0:1], dxyshape) u_dxy = tf.reshape(puv[:Ndxy,1:2], dxyshape) v_dxy = tf.reshape(puv[:Ndxy,2:3], dxyshape) p_hx = tf.reshape(puv[Ndxy:(Ndxy+Nhx),0:1], hxshape) u_hx = tf.reshape(puv[Ndxy:(Ndxy+Nhx),1:2], hxshape) v_hx = tf.reshape(puv[Ndxy:(Ndxy+Nhx),2:3], hxshape) p_hy = tf.reshape(puv[(Ndxy+Nhx):,0:1], hyshape) u_hy = tf.reshape(puv[(Ndxy+Nhx):,1:2], hyshape) v_hy = tf.reshape(puv[(Ndxy+Nhx):,2:3], hyshape) # 计算patch数量（边界无patch） t_patch = dxyshape[0] - 1 #=19 x_patch = dxyshape[1] - 2 #=3 y_patch = dxyshape[2] - 2 #=3 N_patch = t_patch*x_patch*y_patch #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! # 此处为核心问题区域 # 需要创建p_dxy、u_dxy、v_dxy的深拷贝，可被覆盖且支持自动微分（AD）操作追踪 p_proj = tf.identity(p_dxy) #identity函数无法追踪操作，无AD支持 u_proj = tf.TensorArray(u_dxy) #需逐维度创建后赋值，但操作无法被记录 v_proj = tf.Variable(v_dxy) #单epoch可用，但多epoch会触发同名变量创建报错 # 除边界值外将被覆盖，每个*_proj与对应*_dxy形状相同 #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! # 遍历所有patch并处理存储 for ts in range(t_patch): for xs in range(x_patch): for ys in range(y_patch): # patch处理逻辑，输出patch中心的处理后p,u,v值 the_patch = (p_dxy[ts:ts+2, xs:xs+3, ys:ys+3], u_dxy[ts:ts+2, xs:xs+3, ys:ys+3], v_dxy[ts:ts+2, xs:xs+3, ys:ys+3]) p_, u_, v_ = process(the_patch) # 存储处理后的预测值，保留边界值 p_proj[ts+1,xs+1,ys+1].assign(p_) u_proj[ts+1,xs+1,ys+1].assign(u_) v_proj[ts+1,xs+1,ys+1].assign(v_) # 重塑为一维数组 p_proj = tf.reshape(p_proj, [Ndxy]) u_proj = tf.reshape(u_proj, [Ndxy]) v_proj = tf.reshape(v_proj, [Ndxy]) # 组合预测结果 puv_output = tf.stack([p_proj, u_proj, v_proj], axis=1) return puv_output
解决方案
使用**纯张量操作tf.tensor_scatter_nd_update**实现数组更新，该操作返回新张量、保留原张量不变，且完全支持自动微分追踪，避免tf.Variable和tf.TensorArray带来的问题。
修改后的HCP函数核心代码如下：
def HCP(puv, dxyshape, hxshape, hyshape ): '''将预测结果转换为约束patch，处理后得到单个值并输出处理后的结果''' # 提取各类数据的数量 Ndxy = dxyshape[0]*dxyshape[1]*dxyshape[2] Nhx = hxshape[0]*hxshape[1]*hxshape[2] Nhy = hyshape[0]*hyshape[1]*hyshape[2] # 提取p,u,v并重塑为3D数组 p_dxy = tf.reshape(puv[:Ndxy,0:1], dxyshape) u_dxy = tf.reshape(puv[:Ndxy,1:2], dxyshape) v_dxy = tf.reshape(puv[:Ndxy,2:3], dxyshape) p_hx = tf.reshape(puv[Ndxy:(Ndxy+Nhx),0:1], hxshape) u_hx = tf.reshape(puv[Ndxy:(Ndxy+Nhx),1:2], hxshape) v_hx = tf.reshape(puv[Ndxy:(Ndxy+Nhx),2:3], hxshape) p_hy = tf.reshape(puv[(Ndxy+Nhx):,0:1], hyshape) u_hy = tf.reshape(puv[(Ndxy+Nhx):,1:2], hyshape) v_hy = tf.reshape(puv[(Ndxy+Nhx):,2:3], hyshape) # 计算patch数量（边界无patch） t_patch = dxyshape[0] - 1 x_patch = dxyshape[1] - 2 y_patch = dxyshape[2] - 2 N_patch = t_patch*x_patch*y_patch # 初始化新数组，直接以原数组为基础（张量不可变，后续操作生成新张量） p_proj = p_dxy u_proj = u_dxy v_proj = v_dxy # 收集所有需要更新的索引和对应的值 p_updates = [] u_updates = [] v_updates = [] indices = [] # 遍历所有patch并处理存储 for ts in range(t_patch): for xs in range(x_patch): for ys in range(y_patch): # 提取当前patch（基于原数组，不修改原数据） the_patch = (p_dxy[ts:ts+2, xs:xs+3, ys:ys+3], u_dxy[ts:ts+2, xs:xs+3, ys:ys+3], v_dxy[ts:ts+2, xs:xs+3, ys:ys+3]) # patch处理逻辑 p_, u_, v_ = process(the_patch) # 记录更新的索引和值 idx = [ts+1, xs+1, ys+1] indices.append(idx) p_updates.append(p_) u_updates.append(u_) v_updates.append(v_) # 将索引和值转换为张量格式 indices_tensor = tf.convert_to_tensor(indices, dtype=tf.int32) p_updates_tensor = tf.convert_to_tensor(p_updates, dtype=p_dxy.dtype) u_updates_tensor = tf.convert_to_tensor(u_updates, dtype=u_dxy.dtype) v_updates_tensor = tf.convert_to_tensor(v_updates, dtype=v_dxy.dtype) # 使用tf.tensor_scatter_nd_update更新指定位置，生成新张量 p_proj = tf.tensor_scatter_nd_update(p_proj, indices_tensor, p_updates_tensor) u_proj = tf.tensor_scatter_nd_update(u_proj, indices_tensor, u_updates_tensor) v_proj = tf.tensor_scatter_nd_update(v_proj, indices_tensor, v_updates_tensor) # 重塑为一维数组 p_proj = tf.reshape(p_proj, [Ndxy]) u_proj = tf.reshape(u_proj, [Ndxy]) v_proj = tf.reshape(v_proj, [Ndxy]) # 组合预测结果 puv_output = tf.stack([p_proj, u_proj, v_proj], axis=1) return puv_output
方案说明
- 纯张量操作：
tf.tensor_scatter_nd_update是TensorFlow原生张量操作，返回新张量的同时保留原数组不变，完全满足“不覆盖原预测值”的要求。 - 梯度追踪支持：所有操作均为可微分的张量运算，
tf.GradientTape可完整追踪梯度，不会出现梯度为None的情况。 - 避免变量问题：无需创建
tf.Variable，彻底解决tf.function中变量单例的报错问题。
内容的提问来源于stack exchange，提问作者miranda
相关产品推荐
相关产品推荐
