torch.autograd.Function

一昂Yang

于 2020-01-09 15:44:17 发布

阅读量3.5k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： RefineDet pytorch学习

本文链接：https://blog.youkuaiyun.com/shaopeng568/article/details/103909494

pytorch学习同时被 2 个专栏收录

7 篇文章

订阅专栏

RefineDet

1 篇文章

订阅专栏

本文详细介绍了PyTorch中自定义Layer的两种方法：继承torch.nn.Module和torch.autograd.Function。前者适用于大多数情况，而后者则用于自定义求导规则，适合处理不可导操作或PyTorch内置未提供的函数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pytorch自定义layer有两种方式：

方式一：通过继承torch.nn.Module类来实现拓展。只需重新实现__init__和forward函数。

方式二：通过继承torch.autograd.Function，除了要实现__init__和forward函数，还要实现backward函数（就是自定义求导规则）。

方式一看着更简单一点，torch.nn.Module不香么？为毛要用方式二。因为当我们自定义的函数torch.nn.functioanl里没有的时候，或者一些操作不可导，就需要自己定义求导方式，也就是所谓的Extending torch.autograd

官方的示例是这个样子的： Class torch.autograd.Function

>>> class Exp(Function):
>>>
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result

详细学习参见这个博客讲的很好https://blog.youkuaiyun.com/qq_27825451/article/details/95189376