扩展PyTorch
添加操作autograd
需要Function
为每个操作实现一个新的子类。回想一下,Function
使用autograd
来计算结果和梯度,并对操作历史进行编码。每个新功能都需要您实现两种方法:
forward()
- 执行操作的代码。如果您指定了默认值,则可以根据需求使用任意参数,其中一些参数可选。这里支持各种Python
对象。Variable
参数在调用之前会被转换Tensor
,并且它们的使用情况将在graph
中注册。请注意,此逻辑不会遍历lists
/dicts
/和其他任何数据的结构,并且只考虑被直接调用的Variables
参数。如果有多个输出你可以返回单个Tensor
或格式的元组。另外,请参阅Function
文档查找只能被forward()
调用的有用方法的说明。backward()
- 计算梯度的公式. 它将被赋予与输出一样多的Variable
参数, 其中的每一个表示对应梯度的输出. 它应该返回与输入一样多的Variable
, 其中的每一个表示都包含其相应输入的梯度. 如果输入不需要计算梯度 (请参阅needs_input_grad
属性),或者是非
Variable对象,则可返回
None类.此外,如果你在
forward()方法中有可选的参数,
则可以返回比输入更多的梯度,只要它们都是None
类型即可.
现在,为了更方便使用这些自定义操作,推荐使用apply
方法:
我们下面给出一个由非变量参数进行参数化的函数的例子:
你可能想检测你刚刚实现的backward
方法是否正确的计算了梯度。你可以使用小的有限差分法(Finite Difference
)进行数值估计。
添加操作的功能版本已经在上面的章节中已经介绍了。
增加一个Module
。
由于nn
大量使用autograd
。所以, 添加一个新的Module类需要实现一个Function
类, 它会执行对应的操作并且计算梯度。我们只需要很少的代码就可以实现上面Linear
模块的功能。现在,我们需要实现两个函数:
forward()
- 实例化Function
并使用它来执行操作。它与上面显示的functional wrapper
非常相似。
下面是实现Linear
模块的方式: