PyTorch实现含输入梯度计算的前馈神经网络损失不下降问题咨询

• 基于PyTorch实现前馈神经网络过程中，出现损失函数始终无法下降的问题。经过前期测试排查：如果微调网络结构，让网络针对每个输入样本输出2维向量直接作为pred，整个训练流程运行完全正常
• 希望定位当前pred定义逻辑存在的问题，找到让损失可以正常下降的实现方案

import torch
import numpy as np
from torch import nn

dt = 0.1

class Neural_Network(nn.Module):
    def __init__(self, ):
        super(Neural_Network, self).__init__()
    
        self.l1 = nn.Linear(2,300)
        self.nl = nn.Tanh()
        self.l2 = nn.Linear(300,1)
    
    
    def forward(self, X):
        z = self.l1(X)
        z = self.nl(z)
        o = self.l2(z)
        return o



N = 1000
X = torch.rand(N,2,requires_grad=True)
y = torch.rand(N,1)
NN = Neural_Network()

criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.Adam(NN.parameters(), lr=1e-5)

epochs = 200

for i in range(epochs):  # trains the NN 1,000 times

    HH = torch.mean(NN(X))
    gradH =  torch.autograd.grad(HH, X)[0]
    XH= torch.cat((gradH[:,1].unsqueeze(0),-gradH[:,0].unsqueeze(0)),dim=0).t()
    pred = X + dt*XH

    #Optimize and improve the weights
    loss = criterion(pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()


    print (" Loss: ", loss.detach().numpy())  # mean sum squared loss

• 测试阶段使用的X、y为简化设置的随机张量，并不满足让损失趋近于0的分布，设置目的仅为简化示例，验证训练流程能否让损失正常下降，后续这套架构会应用到符合对应模型假设的真实数据集上
• 这套架构的核心目标是借助神经网络拟合矢量场的哈密顿量，训练数据仅包含离散的轨迹观测值，即采样点从x(t) -> x(t+Δt)的状态更新数据：
  • 输入张量X存储所有t时刻的采样点x(t)
  • 标签y存储对应的t+Δt时刻的状态值x(t+Δt)
  • 网络本身用于拟合哈密顿函数H(x)，不直接输出状态预测
• 代码中各计算步骤的对应逻辑：
  • XH为拟合得到的哈密顿量对应的哈密顿矢量场
  • pred = X + dt*XH的计算是采用前向欧拉法实现的单步时间推进
• 抽象后的核心待解决问题：如何在损失函数中合理引入网络针对输入的梯度计算项，保障训练流程可以正常收敛

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