在 PyTorch 函数中使用 PyTorch 模块的正确方法是什么？

Question

我们有一个自定义的torch.autograd.Functionz(x, t)，它以一种不适合直接自动微分的方式计算输出y，并计算了运算相对于其输入x和t的雅可比行列式，所以我们可以实现backward方法。

但是，该操作涉及对神经网络进行多次内部调用，我们现在已将其实现为 torch.nn.Linear 对象的堆栈，包装在 net 中，torch.nn.Module。在数学上，这些是由t 参数化的。

有什么方法可以让net 本身成为z 的forward 方法的输入？然后，我们将从我们的 backward 返回上游梯度 Dy 和参数 Jacobia dydt_i 的产品列表，其中每个参数 ti 是 net 的子级（除了 @987654340 @，虽然x是数据，不需要梯度累积）。

或者我们真的需要采用t（实际上是一个单独的t_i的列表），并在z.forward内部重构Linear中所有Linear层的动作net？

Answer 1

我猜你可以创建一个自定义仿函数来继承 torch.autograd.Function 并使 forward 和 backward 方法非静态（即删除 this example 中的 @staticmethod 以便 net 可以是一个属性你的仿函数。看起来像

class MyFunctor(torch.nn.autograd.Function):
    def __init(net):
         self.net = net
    
     def forward(ctx, x, t):
         #store x and t in ctx in the way you find useful
         # not sure how t is involved here
         return self.net(x) 

     def backward(ctx, grad):
         # do your backward stuff

net = nn.Sequential(nn.Linear(...), ...)
z = MyFunctor(net)
y = z(x, t)

这将产生一个警告，表明您正在使用一种已弃用的传统方法来创建 autograd 函数（由于非静态方法），并且在反向传播后，您需要格外小心将 net 中的梯度归零。所以不是很方便，但我不知道有什么更好的方法来拥有一个有状态的 autograd 函数。