numpy - 定义我自己的梯度函数供 pytorch 使用

Question

我想手动输入 pytorch 梯度。在我的实际问题中，我有自己的不使用张量的伴随函数。有什么方法可以定义我自己的梯度函数供 pytorch 在优化期间使用？

import numpy as np
import torch

# define rosenbrock function and gradient
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])

# create stochastic rosenbrock function and gradient
# (the crude analogy is that I have predefined stochastic
#  forward and backward functions)
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)

# here, closure is fed f_rand to compute the gradient.
# I need to feed closer the gradient directly from 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>)
  # ( this is the right answer, E[f(1, 1)] = 0 )

我试过定义一个自定义函数，但我无法让它工作。这是我迄今为止最好的尝试:

import numpy as np
import torch

# define rosenbrock function and gradient

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])

# create stochastic rosenbrock function and gradient
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:不推荐使用具有非静态转发方法的旧版 autograd 函数。请使用带有静态转发方法的新型 autograd 函数。 (示例:https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function)

Answer 1

不太确定这是否正是您想要的，但是方法调用 loss.backward() 通过 pytorch 的计算图计算梯度并将梯度值存储在权重张量本身中(在您的情况下是在 x_tensor 中)。这些梯度可以通过 x_tensor.grad 访问。但是，如果您不想使用 pytorch 的梯度计算方法使用 loss.backward()，那么您可以手动将梯度输入张量的 .grad 属性，如下所示:

with torch.no_grad():
    def closure():
       optimizer.zero_grad()
       loss = f_rand(x_tensor)
       x_tensor.grad = torch.from_numpy(jac_rand(x_tensor))
       return loss