Pytorch autograd：使一个参数的梯度成为另一个参数的函数

Question

在Pytorch中，如何使参数的梯度成为函数本身？

这是一个简单的代码片段：

import torch

def fun(q):

    def result(w):
        l = w * q
        l.backward()
        return w.grad

    return result

w = torch.tensor((2.), requires_grad=True)
q = torch.tensor((3.), requires_grad=True)

f = fun(q)

print(f(w))

在上面的代码中，如何使 f(w) 相对于 q 具有梯度？

编辑：根据接受的答案，我能够编写有效的代码。本质上，我在 2 个优化步骤之间交替进行。对于 dim == 1 它有效，对于 dim == 2 它无效。我收到错误“RuntimeError：第二次尝试向后遍历图形，但保存的中间结果已被释放。第一次向后调用时指定 retain_graph=True。”

import torch 

class f_class():
    def __init__(self, dim):    
        self.dim = dim
        if self.dim == 1:      
            self.w = torch.tensor((3.), requires_grad=True)
        elif self.dim == 2:            
            self.w = [torch.tensor((3.), requires_grad=True), torch.tensor((5.), requires_grad=True)]
        else:
            raise ValueError("dim 1 or 2")
        
    def forward(self, x):
        if self.dim == 1:      
            return torch.mul(self.w, x)
        elif self.dim == 2:  
            return torch.mul(torch.mul(self.w[0], self.w[1]), x)            
           
    def set_w(self, w):
        self.w = w
        
    def get_w(self):
        return self.w
    
class g_class():
    def __init__(self):                
        self.q = torch.tensor((4.), requires_grad=True)
        
    def forward(self, f):
        return torch.mul(self.q, f)
    
    def set_q(self, q):
        self.q = q
        
    def get_q(self):
        return self.q
        
def w_new(f, g, dim):  
    
    loss_g = g.forward(f.forward(xd))
    
    if dim == 1:
        grads = torch.autograd.grad(loss_g, f.get_w(), create_graph=True, only_inputs=True)[0]
           
        temp = f.get_w().detach() + grads        
    else:                
        grads = torch.autograd.grad(loss_g, f.get_w(), create_graph=True, only_inputs=True)
        
        temp = [wi.detach() + gi for wi, gi in zip(f.get_w(), grads)] 

    return temp    

def q_new(f, g):          

    loss_f = 2 * f.forward(xd) 
       
    loss_f.backward()

    temp = g.get_q().detach() + g.get_q().grad
    
    temp.requires_grad = True
    
    return temp

dim = 1

xd = torch.tensor((2.))
        
f = f_class(dim)
g = g_class()

for _ in range(3):
    
    print(f.get_w(), g.get_q())
    
    wnew = w_new(f, g, dim)
    
    f.set_w(wnew)
    print(f.get_w(), g.get_q())

    qnew = q_new(f, g)
    
    g.set_q(qnew)
    
print(f.get_w(), g.get_q())

Answer 1

在计算梯度时，如果要为梯度本身构建计算图，需要指定create_graph=True自动梯度。

代码中的潜在错误来源是在 f 中使用 Tensor.backward。这里的问题是 w.grad 和 q.grad 将填充 l 的梯度。这意味着当你调用 f(w).backward() 时，f 和 l 的梯度将被添加到 w.grad 和 q.grad。实际上，您最终会得到 w.grad 等于 dl/dw + df/dw 并且 q.grad 也类似。解决这个问题的一种方法是在 f(w) 之后但在 .backward() 之前将梯度归零。更好的方法是在 f 中使用 torch.autograd.grad。使用后一种方法，w和q的grad属性在调用f时不会填充，只有在调用.backward()时才会填充。这为训练过程中的梯度积累留出了空间。

import torch

def fun(q):
    def result(w):
        l = w * q 
        return torch.autograd.grad(l, w, only_inputs=True, retain_graph=True)[0]
    return result


w = torch.tensor((2.), requires_grad=True)
q = torch.tensor((3.), requires_grad=True)

f = fun(q)

f(w).backward()

print('w.grad:', w.grad)
print('q.grad:', q.grad)

结果是

w.grad: None
q.grad: tensor(1.)

请注意 w.grad 未填充。这是因为 f(w) = dl/dw = q 不是 w 的函数，因此 w 不是计算图的一部分。如果您使用标准的 pytorch 优化器，这很好，因为 None 梯度被隐式假定为零。

如果 l 是 w 的非线性函数，那么 w.grad 将在 f(w).backward() 之后填充。例如

import torch

def fun(q):
    def result(w):
        # now dl/dw = 2 * w * q
        l = w**2 * q
        return torch.autograd.grad(l, w, only_inputs=True, create_graph=True)[0]
    return result

w = torch.tensor((2.), requires_grad=True)
q = torch.tensor((3.), requires_grad=True)

f = fun(q)

f(w).backward()

print('w.grad:', w.grad)
print('q.grad:', q.grad)

结果是

w.grad: tensor(6.)
q.grad: tensor(4.)

Pytorch autograd：使一个参数的梯度成为另一个参数的函数

Pytorch autograd: Make gradient of a parameter a function of another parameter

pytorch

autograd