如何为PyTorch自定义梯度函数以手动提供梯度？

问题描述 #

希望手动向PyTorch提供梯度，实际场景中有不使用张量的自定义伴随函数，能否定义专属梯度函数供PyTorch优化过程使用？

初始示例代码 #

import numpy as np
import torch

# 定义Rosenbrock函数及其梯度
x0 = np.array([0.1, 0.1])
a = 1
b = 5
def f(x):
   return (a - x[0]) ** 2 + b * (x[1] - x[0] ** 2) ** 2

def jac(x):
   dx1 = -2 * a + 4 * b * x[0] ** 3 - 4 * b * x[0] * x[1] + 2 * x[0]
   dx2 = 2 * b * (x[1] - x[0] ** 2)
   return np.array([dx1, dx2])

# 构造带随机噪声的Rosenbrock函数及其梯度
def f_rand(x):
   return f(x) * np.random.uniform(0.5, 1.5)

def jac_rand(x): return jac(x) * np.random.uniform(0.5, 1.5)


x_tensor = torch.tensor(x0, requires_grad=False)
optimizer = torch.optim.Adam([x_tensor], lr=0.1)

# 闭包中通过f_rand计算梯度，需求是直接传入jac_rand的梯度
def closure():
   optimizer.zero_grad()
   loss = f_rand(x_tensor)
   loss.backward() # 实际期望是jac_rand(x)
   return loss

for ii in range(200):
   optimizer.step(closure) 

print(x_tensor, f(x_tensor))
# 输出：tensor([1.0000, 1.0000], dtype=torch.float64, requires_grad=True) tensor(4.5799e-09, dtype=torch.float64, grad_fn=<AddBackward0>)
# 正确结果，E[f(1, 1)] = 0

自定义函数尝试代码及报错 #

尝试通过torch.autograd.Function定义自定义梯度函数，但运行报错：

import numpy as np
import torch

# 定义Rosenbrock函数及其梯度
x0 = np.array([0.1, 0.1])
a = 1
b = 5
def f(x):
   return (a - x[0]) ** 2 + b * (x[1] - x[0] ** 2) ** 2

def jac(x):
   dx1 = -2 * a + 4 * b * x[0] ** 3 - 4 * b * x[0] * x[1] + 2 * x[0]
   dx2 = 2 * b * (x[1] - x[0] ** 2)
   return np.array([dx1, dx2])

# 构造带随机噪声的Rosenbrock函数及其梯度
def f_rand(x):
   return f(x) * np.random.uniform(0.5, 1.5)

def jac_rand(x): return jac(x) * np.random.uniform(0.5, 1.5)

class custom_function(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
       ctx.save_for_backward(input)
       return f_rand(input)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * g_rand(input)

x_tensor = torch.tensor(x0, requires_grad=False)
optimizer = torch.optim.Adam([x_tensor], lr=0.1)

for ii in range(200):
   print('x_tensor ', x_tensor)
   optimizer.step(custom_function())

print(x_tensor, f(x_tensor))

报错信息：

RuntimeError: 旧版自动求导函数的非静态forward方法已被弃用，请使用新版带静态forward方法的自动求导函数。

修正后的解决方案 #

代码存在多个问题，逐一修正后可正常运行：

1. 笔误修正：backward方法中g_rand应为jac_rand
2. 启用梯度追踪：x_tensor需设置requires_grad=True，否则无法累积梯度
3. 张量类型统一：forward返回值需从numpy数组转为PyTorch张量；backward中jac_rand的输出也要转成张量，保证与grad_output的设备、数据类型一致
4. 优化器闭包正确使用：optimizer.step()需要传入闭包函数，而非直接传入函数实例

修正后的完整代码：

import numpy as np
import torch

# 定义Rosenbrock函数及其梯度
x0 = np.array([0.1, 0.1])
a = 1
b = 5
def f(x):
   return (a - x[0]) ** 2 + b * (x[1] - x[0] ** 2) ** 2

def jac(x):
   dx1 = -2 * a + 4 * b * x[0] ** 3 - 4 * b * x[0] * x[1] + 2 * x[0]
   dx2 = 2 * b * (x[1] - x[0] ** 2)
   return np.array([dx1, dx2])

# 构造带随机噪声的Rosenbrock函数及其梯度
def f_rand(x):
   return f(x) * np.random.uniform(0.5, 1.5)

def jac_rand(x): return jac(x) * np.random.uniform(0.5, 1.5)

class custom_function(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
       ctx.save_for_backward(input)
       # 将numpy结果转为tensor
       return torch.tensor(f_rand(input.numpy()), dtype=input.dtype, device=input.device)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        # 计算自定义梯度并转为tensor，匹配grad_output的设备和类型
        grad = torch.tensor(jac_rand(input.numpy()), dtype=grad_output.dtype, device=grad_output.device)
        return grad_output * grad

x_tensor = torch.tensor(x0, requires_grad=True)
optimizer = torch.optim.Adam([x_tensor], lr=0.1)

def closure():
    optimizer.zero_grad()
    # 调用自定义函数，需用apply方法
    loss = custom_function.apply(x_tensor)
    loss.backward()
    return loss

for ii in range(200):
    optimizer.step(closure) 

print(x_tensor, f(x_tensor.numpy()))

运行后可得到与初始示例一致的收敛结果，成功通过自定义梯度函数完成PyTorch优化流程。

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