译者:PEGASUS1993
本章中,将要介绍使用我们的C库如何扩展torch.nn
,torch.autograd
和编写自定义的C
扩展工具。
扩展torch.autograd
添加操作autograd
需要Function
为每个操作实现一个新的子类。回想一下,Function
使用autograd
来计算结果和梯度,并对操作历史进行编码。每个新功能都需要您实现两种方法:
-
forward()
- 执行操作的代码。如果您指定了默认值,则可以根据需求使用任意参数,其中一些参数可选。这里支持各种Python
对象。Variable
参数在调用之前会被转换Tensor
,并且它们的使用情况将在graph
中注册。请注意,此逻辑不会遍历lists
/dicts
/和其他任何数据的结构,并且只考虑被直接调用的Variables
参数。如果有多个输出你可以返回单个Tensor
或Tensor
格式的元组。另外,请参阅Function
文档查找只能被forward()
调用的有用方法的说明。 -
backward()
- 计算梯度的公式. 它将被赋予与输出一样多的Variable
参数, 其中的每一个表示对应梯度的输出. 它应该返回与输入一样多的Variable
, 其中的每一个表示都包含其相应输入的梯度. 如果输入不需要计算梯度 (请参阅needs_input_grad
属性),或者是非Variable
对象,则可返回None
类.此外,如果你在forward()
方法中有可选的参数,则可以返回比输入更多的梯度,只要它们都是None
类型即可.
你可以从下面的代码看到torch.nn
模块的Linear
函数, 以及注解
# Inherit from Function
class Linear(Function):
# bias is an optional argument
def forward(self, input, weight, bias=None):
self.save_for_backward(input, weight, bias)
output = input.mm(weight.t())
if bias is not None:
output = bias.unsqueeze(0).expand_as(output)
return output
# This function has only a single output, so it gets only one gradient
def backward(self, grad_output):
# This is a pattern that is very convenient - at the top of backward
# unpack saved_tensors and initialize all gradients w.r.t. inputs to
# None. Thanks to the fact that additional trailing Nones are
# ignored, the return statement is simple even when the function has
# optional inputs.
input, weight, bias = self.saved_tensors
阅读全文/改进本文