DeepXDE含可训练变量初始化时PyTorch双重反向传播错误排查

• 检查可训练变量x0的定义与注册

  • 必须用torch.nn.Parameter()定义x0，并确保它被正确注册到模型的参数集合中。如果是自定义模块，要将x0作为模块属性；如果是直接在DeepXDE框架内使用，需通过DeepXDE提供的参数注册方式添加，避免因张量未被优化器追踪引发梯度状态混乱。禁止使用普通torch张量替代Parameter类型。

• 排查初始条件函数的张量运算路径

  • 彻底替换原numpy分布代码，确保fun_init中所有运算基于PyTorch张量完成，杜绝numpy与torch张量混合计算——混合运算会导致部分张量无梯度信息，触发反向传播异常。
  • 检查高斯分布生成逻辑，比如torch.distributions.Normal(x0, std).sample()这类操作，不要用detach()或no_grad()包裹需要参与梯度传播的部分，确保计算图完整。

• 验证损失计算的梯度流向

  • 在训练循环的total_loss.backward()前手动执行optimizer.zero_grad()，强制清除上一轮迭代的梯度。虽然DeepXDE默认处理梯度清零，但自定义可训练变量时可能存在覆盖不全的情况，避免梯度累积引发双重反向传播冲突。
  • 打印x0的梯度状态：第一次反向传播后查看x0.grad是否存在，第二次迭代前确认梯度已被清零，通过梯度状态判断反向传播流程是否正常。

• 禁用自动混合精度（AMP）测试

  • 暂时关闭DeepXDE或PyTorch的AMP设置，部分场景下AMP会导致梯度张量类型不匹配，尤其是自定义可训练变量时，以此排查精度相关的反向传播问题。

• 简化模型定位问题根源

  • 先只保留初始条件损失项，验证是否仍触发错误；再逐步添加其他损失项，定位引发问题的损失分支。
  • 使用极小数据集和简单网络结构测试，排除复杂模型结构带来的干扰，快速缩小问题范围。

• 检查版本兼容性

  • 尝试降级PyTorch到2.1.x版本，或升级DeepXDE到最新稳定版。新版本PyTorch的自动微分逻辑变化可能与旧版DeepXDE存在兼容性问题，版本匹配后可能解决错误。

内容的提问来源于stack exchange，提问作者Miles N.