PyTorch LBFGS优化器下参数乘法被识别为原地操作报错如何解决

根本原因 #

你遇到的报错和r*theta乘法操作本身无关，是LBFGS优化器的使用逻辑错误导致的：

• LBFGS优化器在单次step执行过程中会多次调用传入的closure函数，用于评估不同参数点的损失与梯度，且参数更新为原地操作。
• 你当前代码中将损失Wass的计算逻辑放在了closure外部，第一次调用closure完成反向传播、参数更新后，最初用于计算Wass的r、theta值已经被原地修改覆盖。
• 加法操作的反向传播不需要依赖前向传播时的输入原始值（梯度恒为1），因此即使参数被修改也不会触发报错；但乘法操作的反向传播需要用到前向传播时存储的r、theta原始值，旧值被覆盖后就会触发原地操作修改梯度依赖变量的报错。

修复方案 #

把所有前向传播、损失计算的逻辑全部移到closure函数内部，保证每次调用closure时，都基于当前最新的参数重新构建计算图，不会依赖已经被修改的旧参数值。

修改后的核心代码段 #

def learn_U_param(Zobs, ngrad, params, r_guess=0., theta_guess=0., lambd=1.):
    Npts = params[0]
    positions = torch.tensor(np.arange(0, 1, 1 / Npts) + 1 / 2 / Npts).reshape((Npts, 1))
    lab = torch.tensor(np.arange(0, Npts))
    # 初始化可训练参数
    r = torch.tensor(float(r_guess)).to(device)
    r.requires_grad = True
    theta = torch.tensor(float(theta_guess)).to(device)
    theta.requires_grad = True
    r_hist = [r.item()]
    theta_hist = [theta.item()]
    loss_hist = []
    optimizer = optim.LBFGS([r, theta])
    for i in range(ngrad):
        def closure():
            optimizer.zero_grad()
            # 所有前向、损失计算都移到closure内部
            field = Wrong_U_param(r, theta, positions)
            Z = model(field)
            Loss = SamplesLoss(loss="sinkhorn", p=2, blur=.05)
            Wass = Loss(lab, Z, positions, lab, Zobs, positions)
            Wass.backward()
            return Wass
        optimizer.step(closure)
        # 每次迭代结束后记录当前参数和损失
        with torch.no_grad():
            field = Wrong_U_param(r, theta, positions)
            Z = model(field)
            Loss = SamplesLoss(loss="sinkhorn", p=2, blur=.05)
            Wass = Loss(lab, Z, positions, lab, Zobs, positions)
            loss_hist.append(Wass.item())
        r_hist.append(r.item())
        theta_hist.append(theta.item())
    return r_hist, theta_hist, loss_hist

额外说明 #

• 不需要修改r*theta的实现，不需要用item()或者其他破坏计算图的操作，调整计算逻辑后原生乘法即可正常计算梯度。
• 修复后代码可以正常优化r和theta的值，最终会收敛到设置的真值r=2、theta=2附近。

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