扩展PyTorch

1. 扩展 torch.autograd
2. 扩展 torch.nn
3. 编写自定义C扩展

本篇文章中包含如何扩展torch.nn，torch.autograd和使用C库编写自定义C扩展。

扩展torch.autograd

添加操作autograd需要Function为每个操作实现一个新的 子类。回想一下，Functions是autograd用于计算结果和渐变，以及编码操作历史。每个新功能都需要您实现2种方法：

1. forward() - 执行操作的代码。如果您指定了默认值，则可以使用任意数量的参数，其中一些参数可选。在这里接受各种Python对象。 Variable参数将Tensor在调用之前转换为s，并且它们的使用将在图表中注册。请注意，此逻辑不会遍历列表/ dicts /任何其他数据结构，并且只会考虑作为调用的直接参数的变量。您可以返回一个单一的 Tensor输出，或tuple中Tensor如果有多个输出秒。另外，请参阅文档Function来查找只能从中调用的有用方法的说明forward()。
2. backward() - 梯度公式。将给出与Variable输出一样多的参数，每个参数表示输出的渐变wrt。它应该返回尽可能多 Variable的输入，其中每个包含渐变wrt其相应的输入。如果您的输入不需要渐变（参见needs_input_grad）或非Variable 对象，则可以返回None。此外，如果您有可选参数，则forward()可以返回比输入更多的渐变，只要它们全部None。

面你可以找到一个Linear函数的代码torch.nn，并附加其他注释：

# 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
        grad_input = grad_weight = grad_bias = None

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        if self.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if self.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and self.needs_input_grad[2]:
            grad_bias = grad_output.sum(0).squeeze(0)

        return grad_input, grad_weight, grad_bias

现在，为了更容易使用这些自定义操作，我们建议使用别名的apply方法：

linear = Linear.aply

在这里，我们给一个函数，参数化non-Variable的另外的一个例子：

def linear(input, weight, bias=None):
    # First braces create a Function object. Any arguments given here
    # will be passed to __init__. Second braces will invoke the __call__
    # operator, that will then use forward() to compute the result and
    # return it.
    return Linear()(input, weight, bias)

你可能想知道你刚刚实现的 backward方法是否正确的计算了梯度。你可以使用小的有限的差分进行数值估计。

from torch.autograd import gradcheck

# gradchek takes a tuple of tensor as input, check if your gradient
# evaluated with these tensors are close enough to numerical
# approximations and returns True if they all verify this condition.
input = (Variable(torch.randn(20,20).double(), requires_grad=True),)
test = gradcheck.gradcheck(Linear(), input, eps=1e-6, atol=1e-4)
print(test)

扩展 torch.nn

nn 包含两种接口 - modules和他们的functional版本。通过这两个接口，你都可以扩展nn。但是我们建议，在扩展layer的时候，使用modules， 因为modules保存着参数和buffer。如果不需要参数的话，那么建议使用functional(激活函数，pooling等)。

增加一个operation的 functional版本已经在上面一节介绍完毕。

增加一个模块(module)。

由于nn大量使用autograd。所以，添加一个新Module需要实现一个用来执行计算梯度的Function。从现在开始，假定我们想要实现一个Linear模块，上面我们已经实现了一个Linear Funciton。 只需要很少的代码就可以完成这个工作。 现在，我们需要实现两个方法：

1. __init__ (optional) - 接收诸如内核大小，功能数量等参数，并初始化参数和缓冲区。
2. forward() - 实例化Function并使用它来执行操作。它与上面显示的功能包装非常相似。

下面是一个Linear模块实现的方式：

class Linear(nn.Module):
    def __init__(self, input_features, output_features, bias=True):
        self.input_features = input_features
        self.output_features = output_features

        # nn.Parameter is a special kind of Variable, that will get
        # automatically registered as Module's parameter once it's assigned
        # as an attribute. Parameters and buffers need to be registered, or
        # they won't appear in .parameters() (doesn't apply to buffers), and
        # won't be converted when e.g. .cuda() is called. You can use
        # .register_buffer() to register buffers.
        # nn.Parameters can never be volatile and, different than Variables,
        # they require gradients by default.
        self.weight = nn.Parameter(torch.Tensor(input_features, output_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(output_features))
        else:
            # You should always register all possible parameters, but the
            # optional ones can be None if you want.
            self.register_parameter('bias', None)

        # Not a very smart way to initialize weights
        self.weight.data.uniform_(-0.1, 0.1)
        if bias is not None:
            self.bias.data.uniform_(-0.1, 0.1)

    def forward(self, input):
        # See the autograd section for explanation of what happens here.
        return Linear()(input, self.weight, self.bias)
        #注意这个Linear是之前实现过的Linear

编写自定义C扩展

即将发布。不过现在你可以在 GitHub找到一个例子 。