TensorFlow自定义函数梯度问题：地球物理反演正演算子适配训练

我想用深度学习解决地球物理经典反演问题：网络输入为测量值d，输出为电阻率模型m。常规损失用MSE(m, m_pred)，但因为是物理问题，需要加入数据响应正则项d_pred=F[m_pred]（F为正演模拟算子），损失函数为：
$$\mathcal{L} = \lambda \cdot \text{MSE}(m, \hat{m}) + (1-\lambda) \cdot \text{MSE}(d, F(\hat{m}))$$

当前正演算子实现代码如下：

import math
import cmath
import time
from scipy import constants
import numpy as np  # 补充原代码缺失的numpy导入

mu = constants.mu_0; #Magnetic Permeability (H/m)

# forward modelling operator F[m] = d
def MTforwardModel(resistivities,thicknesses,frequencies):
  n = len(resistivities);
  apparentResistivity=[]
  phase=[]
  for frequency in frequencies:
      w =  2*math.pi*frequency;
      impedances = list(range(n));
      #compute basement impedance
      impedances[n-1] = np.sqrt(w*mu*resistivities[n-1]*1j);

      for j in range(n-2,-1,-1):
          resistivity = resistivities[j];
          thickness = thicknesses[j];

          # Step 2. Iterate from bottom layer to top(not the basement)
          # Step 2.1 Calculate the intrinsic impedance of current layer
          dj = np.sqrt((w * mu * (1.0/resistivity))*1j);
          wj = dj * resistivity;
          # Step 2.2 Calculate Exponential factor from intrinsic impedance
          ej = np.exp(-2*thickness*dj);

          # Step 2.3 Calculate reflection coeficient using current layer
          #          intrinsic impedance and the below layer impedance
          belowImpedance = impedances[j + 1];
          rj = (wj - belowImpedance)/(wj + belowImpedance);
          re = rj*ej;
          Zj = wj * ((1 - re)/(1 + re));
          impedances[j] = Zj;

      # Step 3. Compute apparent resistivity from top layer impedance
      Z = impedances[0];
      absZ = abs(Z);
      #apparentResistivity.append((absZ * absZ)/(mu * w))
      #phase.append(math.atan2(Z.imag, Z.real))
      phase.append(absZ)
  return np.array(phase)#(np.array(apparentResistivity),np.array(phase))

注：resistivities是网络输出，可复现问题。仅使用正演模拟项（损失函数中的右侧项）时，出现错误：ValueError: No gradients provided for any variable，推测是正演算子内部导致网络与输出断开，无法计算梯度。

现在有两个问题：

1. 如何编写正演算子以适配损失函数及神经网络训练时的梯度计算？
2. 若正演模拟项当前无法计算梯度但能计算数值，将其与MSE项结合后，损失是否会参与反向传播及权重更新？

问题解答 #

1. 适配梯度计算的正演算子编写方案 #

当前正演算子无法计算梯度的核心原因是混用了numpy和原生Python的非自动微分操作，且未使用深度学习框架（如TensorFlow/PyTorch）的张量运算。要让正演算子支持反向传播，需做以下修改：

• 替换所有numpy操作为框架原生张量运算：比如用PyTorch的torch.sqrt()、torch.exp()、torch.abs()替代numpy对应函数；用框架的复数运算替代cmath。
• 避免Python原生列表的循环赋值：将impedances改为框架的张量数组，用索引操作替代列表赋值，确保整个计算过程被框架的计算图追踪。
• 确保所有输入输出都是框架张量：网络输出的resistivities必须是框架张量，thicknesses、frequencies也需转为张量（若为固定参数，可设置为不需要梯度的常量张量）。

以PyTorch为例，修改后的正演算子大致结构如下：

import torch
from scipy import constants

mu = constants.mu_0

def MTforwardModel_torch(resistivities, thicknesses, frequencies):
    n = resistivities.shape[0]
    # 将输入转为框架张量，设置合适的数据类型
    frequencies = torch.tensor(frequencies, dtype=torch.complex64)
    thicknesses = torch.tensor(thicknesses, dtype=torch.complex64)
    w = 2 * torch.pi * frequencies
    # 初始化阻抗张量，维度匹配频率数和层数
    impedances = torch.zeros((len(frequencies), n), dtype=torch.complex64)
    
    # 计算基底阻抗
    impedances[:, -1] = torch.sqrt(w * mu * resistivities[-1] * 1j)
    
    for j in range(n-2, -1, -1):
        rho = resistivities[j]
        h = thicknesses[j]
        # 计算本征阻抗
        dj = torch.sqrt((w * mu * (1.0 / rho)) * 1j)
        wj = dj * rho
        # 指数项
        ej = torch.exp(-2 * h * dj)
        # 反射系数
        below_impedance = impedances[:, j+1]
        rj = (wj - below_impedance) / (wj + below_impedance)
        re = rj * ej
        Zj = wj * ((1 - re) / (1 + re))
        impedances[:, j] = Zj
    
    Z = impedances[:, 0]
    absZ = torch.abs(Z)
    return absZ

修改后整个计算过程都在PyTorch的计算图中，能自动追踪梯度，解决"No gradients provided"的问题。

2. 非梯度项与MSE结合后的反向传播情况 #

如果正演模拟项无法计算梯度，仅能输出数值，那么：

• 当把它与MSE项相加作为总损失时，只有MSE项会参与反向传播，正演模拟项的数值相当于一个固定的偏移量，不会对网络权重更新产生任何影响。
• 因为自动微分框架只会对可追踪梯度的张量运算计算梯度，非张量的数值或断开计算图的操作不会贡献梯度，总损失的梯度就等于MSE项的梯度。

内容的提问来源于stack exchange，提问作者Paul Goyes