PyTorch自定义反向传播函数是否需修改？技术咨询

你的判断是正确的，不需要修改原模型的反向传播函数，也不需要改动前向传播逻辑，原因如下：

• PyTorch的自动微分机制会链式组合梯度：

  • 你已经为原模型实现了针对参数P的dA/dP和dB/dP，这相当于在计算图中为A、B到P的路径提供了梯度信息。
  • 当你定义f(A,B) = A * g(B)时，PyTorch会自动计算f对A、B的局部梯度：df/dA = g(B)，df/dB = A * g'(B)（只要g(B)本身是用PyTorch原生可微算子实现的）。
  • 之后链式法则会自动把这两部分局部梯度和你自定义的dA/dP、dB/dP结合，得到最终的df/dP = (dA/dP)*g(B) + A*g'(B)*(dB/dP)，完全符合你预期的梯度公式。

• 仅需注意的细节：

  • 确保g(B)的实现基于PyTorch原生可微算子，这样PyTorch能自动计算g'(B)；如果g也是自定义的非AD实现，那你需要为g单独实现反向传播，但这和原模型的反向传播无关。
  • 原模型的反向传播函数只需要保证能正确输出A、B对P的梯度即可，无需针对f做任何修改。

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