Scipy L-BFGS-B转PyTorch LBFGS自定义梯度适配问题

已使用scipy.optimize.minimize的L-BFGS-B算法实现带自定义梯度的优化，现需迁移至PyTorch实现等价逻辑，但遇到梯度形状不匹配的错误。

Scipy原始代码 #

from scipy.optimize import minimize
import numpy as np

res = minimize(calc_cost, x_0, args=const_data, method='L-BFGS-B', jac=calc_grad)

def calc_cost(x, const_data):
    # 计算得到数组"calculation"
    return np.sum(np.square(calculation)) # 返回标量损失

def calc_grad(x, const_data):
    # 计算得到数组"calculation"（此处为目标函数对x的梯度）
    return np.ravel(calculation) # PyTorch版本无需ravel

遇到的问题 #

参考示例编写PyTorch代码时，触发错误：

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([3, 200, 200]) and output[0] has a shape of torch.Size([])

当前PyTorch代码：

optimizer = optim.LBFGS([x_0], history_size=10, max_iter=10, line_search_fn="strong_wolfe")
h_lbfgs = []
for i in range(10):
    optimizer.zero_grad()
    objective = calc_cost(x_0, const_data)
    objective.backward(gradient = calc_gradient(x_0, const_data))
    optimizer.step(lambda: calc_cost(x_0, const_data))
    h_lbfgs.append(objective.item())

已知梯度形状需与目标函数（标量）匹配，但不符合实际需求，需调整代码实现与Scipy版本一致的计算逻辑。

解决方案 #

核心问题是对PyTorch中backward参数和L-BFGS优化器逻辑的误解：

• Scipy的jac参数直接传入目标函数对输入x的梯度（形状与x一致）；
• PyTorch中backward(gradient=...)的参数是上游梯度（输出对自身的梯度），标量损失的上游梯度只能是标量，不能直接传入对x的梯度。

正确做法是手动将自定义计算的梯度赋值给x_0.grad，同时遵循PyTorch L-BFGS对closure函数的要求：

修改后的PyTorch代码 #

import torch
import torch.optim as optim

# 确保x_0是requires_grad=True的张量，示例：
# x_0 = torch.tensor(初始值, dtype=torch.float32, requires_grad=True)

optimizer = optim.LBFGS([x_0], history_size=10, max_iter=10, line_search_fn="strong_wolfe")
h_lbfgs = []

def closure():
    optimizer.zero_grad()
    # 计算标量损失，对应Scipy的calc_cost
    objective = calc_cost(x_0, const_data)
    # 计算目标函数对x的梯度，对应Scipy的calc_grad
    grad = calc_gradient(x_0, const_data)
    # 确保梯度形状与x_0完全一致（对应Scipy的ravel操作，按需调整reshape逻辑）
    x_0.grad = grad.reshape(x_0.shape)
    return objective

for i in range(10):
    # L-BFGS会多次调用closure做线搜索，必须传入closure而非仅损失函数
    obj_val = optimizer.step(closure)
    h_lbfgs.append(obj_val.item())

关键说明 #

1. Closure函数要求：PyTorch的L-BFGS优化器依赖closure函数完成梯度计算的闭环（清零梯度、计算损失、赋值梯度、返回损失），因为算法需要多次评估损失和梯度做线搜索，不能像SGD那样简单调用step。
2. 梯度形状匹配：calc_gradient返回的张量必须与x_0形状严格一致，可通过reshape调整（对应Scipy中的ravel，如果Scipy中x是展平的，PyTorch中也需保证梯度展平后与x_0形状匹配）。
3. 自定义梯度赋值：直接将计算好的目标对x的梯度赋值给x_0.grad，替代backward(gradient=...)的错误用法，这完全等价于Scipy中jac参数的作用。
4. calc_cost/calc_gradient的PyTorch适配：确保这两个函数内部用PyTorch张量操作实现，比如将np.sum(np.square(...))替换为torch.sum(torch.square(...))，保证计算逻辑与Scipy完全一致。

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