PyTorch中autograd.grad()梯度始终为零的问题排查求助

问题描述 #

实现了自定义损失函数，训练简单全连接神经网络时，发现损失函数的梯度始终为零，无法推进模型训练。怀疑问题出在output_grad, _ = torch.autograd.grad(q, (x,y))这一行。训练时，configuration是从坐标(x,y)确定的分布中采样的点。

目标函数代码 #

def objective(p, output):
  x,y = p
  a = minA
  b = minB
  r = 0.1

  XA = 1/2 -1/2 * torch.tanh(100*((x - a[0])**2 + (y - a[1])**2 - (r + 0.02)**2))
  XB = 1/2 -1/2 * torch.tanh(100*((x - b[0])**2 + (y - b[1])**2 - (r + 0.02)**2))

  q = (1-XA)*((1-XB)* output + (XB))

  output_grad, _ = torch.autograd.grad(q, (x,y))
  output_grad.requires_grad_()
  q = output_grad**2

  return q

训练代码 #

model = NN(input_size)
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
for e in range(epochs) :
  for configuration in total:
    print("Train for configuration", configuration)
    # Training pass
    optimizer.zero_grad()

    #output is q~
    output = model(configuration)

    #loss is the objective function we defined
    loss = objective(configuration, output.item())
    loss.backward()
    optimizer.step()

问题根源分析 #

1. 计算图被切断：调用output.item()将模型输出的张量转换为Python数值，彻底切断了损失函数与模型参数之间的计算图连接，导致梯度无法反向传播到模型参数。
2. 梯度计算逻辑错误：torch.autograd.grad(q, (x,y))计算的是q对输入x,y的梯度，但x,y是输入数据而非模型参数，这部分梯度与模型参数无关；同时该函数返回的张量默认不保留计算图，即使手动设置requires_grad_()，也无法和原模型的计算图关联，无法将梯度传递回模型参数。
3. Tanh函数饱和：tanh(100*...)的系数100过大，导致Tanh函数在大部分输入区域处于饱和状态（输出接近±1），此时函数的梯度趋近于零，进一步加剧了梯度消失问题。

解决方案 #

修改后的目标函数 #

def objective(p, output):
    x, y = p
    a = minA
    b = minB
    r = 0.1

    # 降低Tanh的系数，避免函数过早饱和，保留梯度信息
    XA = 1/2 - 1/2 * torch.tanh(10 * ((x - a[0])**2 + (y - a[1])**2 - (r + 0.02)**2))
    XB = 1/2 - 1/2 * torch.tanh(10 * ((x - b[0])**2 + (y - b[1])**2 - (r + 0.02)**2))

    q = (1 - XA) * ((1 - XB) * output + XB)
    
    # 设置create_graph=True，保留梯度的计算图，让梯度可以反向传播到模型参数
    # grad_outputs指定梯度的初始值，这里用ones_like(q)保证梯度维度匹配
    output_grad = torch.autograd.grad(q, (x, y), grad_outputs=torch.ones_like(q), create_graph=True)[0]
    loss = output_grad ** 2

    return loss.sum()  # 将张量损失转换为标量，满足backward()的要求

修改后的训练代码 #

model = NN(input_size)
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
for e in range(epochs):
    for configuration in total:
        print("Train for configuration", configuration)
        optimizer.zero_grad()

        # 保留output的张量形式，不要转换为Python数值，维持计算图连接
        output = model(configuration)

        loss = objective(configuration, output)
        loss.backward()
        optimizer.step()

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