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

Question

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

这是一个简单的代码sn-p：

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 具有梯度？

编辑：根据接受的答案，我能够编写一个有效的代码。本质上，我在两个优化步骤之间交替。对于 dim == 1 有效，对于 dim == 2 无效。我收到错误“RuntimeError: Trying back through the graph a second time, but the saved intermediate results have been freed.请在第一次向后调用时指定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 指定为autograd。

代码中的一个潜在错误来源是在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.)