利用torch.autograd.Function自定义层的forward和backward

当我们在Pytorch中想自定义某一层的梯度计算时,可以利用torch.autograd.Function来封装一个class,此时可以我们可以自己在backward方法中自定求解梯度的方法,也适用于不可导函数的backward计算。 这个函数的源代码可以从如下链接获取: https://pytorch.org/docs/stable/_modules/torch/autograd/function.html 首先给出一个官方提供的demo: class Exp(torch.autograd.Function): @staticmethod def forward(ctx, i