torch.optim.LBFGS()图像角点优化参数无变化及NaN问题排查

初始参数不更新的根源 #

你最初的forward方法返回torch.autograd.Variable(cumulative_loss.data, requires_grad=True)，其中cumulative_loss.data会直接剥离张量的计算图，得到无梯度关联的纯数值张量，后续手动设置requires_grad=True也无法重建梯度传递链路，导致反向传播时梯度无法传递到待优化的self._real_corners参数，优化器自然不会更新参数。

改成return cumulative_loss.requires_grad_()后，梯度链路恢复，参数开始更新，但很快出现NaN，下面是具体原因和修复方案：

参数NaN异常的修复方案 #

1. 损失计算的数值稳定性优化 #

你的损失计算包含torch.sqrt和torch.pow，当loss_map平方和趋近于0时易出现数值下溢，或平方和过大时出现溢出，进而产生NaN。

• 修复：给开根号的项添加极小epsilon，避免除以0或开根号0：

  cumulative_loss = torch.sqrt(torch.sum(torch.pow(loss_map, 2)) / (loss_map.size(dim=2)*loss_map.size(dim=3)) + 1e-8)
  

• 可选：直接使用MSE损失（无需开根号），减少数值波动：

  cumulative_loss = torch.mean(torch.pow(loss_map, 2))
  

2. L-BFGS学习率过高 #

L-BFGS是自适应学习率算法，你设置的lr=1过大，会导致参数更新幅度过大，超出合理范围，进而引发单应性矩阵计算异常产生NaN。

• 修复：降低初始学习率，比如设置为1e-3或1e-2，同时可限制单次迭代步长：

  optimizer = torch.optim.LBFGS(optimize_corners.parameters(), lr=1e-3, max_iter=20, tolerance_grad=1e-7)
  

3. 单应性矩阵计算的鲁棒性不足 #

当self._real_corners坐标偏离合理范围时，find_homography_dlt可能生成奇异矩阵，后续透视变换会产生NaN。

• 修复1：给参数添加范围约束，限制坐标在图像尺寸内：
  在forward方法开头添加：

  img_h, img_w = real_image.shape[2], real_image.shape[3]
  self._real_corners.data = torch.clamp(self._real_corners.data, min=0, max=max(img_h, img_w))
  

• 修复2：替换为鲁棒性更强的RANSAC版本单应性估计：

  homography_matrix, inliers = kornia.geometry.homography.find_homography_ransac(
      synt_corners, self._real_corners, confidence=0.99, max_iterations=100
  )
  

4. Threshold层的梯度突变问题 #

torch.nn.Threshold会导致梯度突变，若阈值设置不合理，易引发损失异常。

• 修复：改用平滑的截断方式，避免梯度突变：

  # 等价于大于_threshold时置0，否则保留原值
  loss_map = torch.min(loss_map, torch.tensor(_threshold, device=loss_map.device))
  

L-BFGS优化器的正确使用方式 #

你的优化循环写法不符合L-BFGS的要求，官方推荐在step的闭包内完成损失计算和梯度反向传播，正确写法：

for it in tqdm(range(n_iters), desc='Fitting corners', leave=False, position=1):
    def closure():
        optimizer.zero_grad()
        loss = optimize_corners(_image, _synt_image, _synt_corners, _threshold)
        loss.backward()
        return loss
    # step会自动执行闭包内的梯度计算，并更新参数
    loss = optimizer.step(closure)
    h_lbfgs.append(loss.item())
    print(h_lbfgs)

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