總說

雖然pytorch可以自動求導，但是有時候一些操作是不可導的，這時候你需要自定義求導方式。也就是所謂的 “Extending torch.autograd”. 官網雖然給了例子，但是很簡單。這裡將會更好的說明。

擴充套件 torch.autograd

class LinearFunction(Function):

    # 必須是staticmethod
    @staticmethod
    # 第一個是ctx，第二個是input，其他是可選引數。
    # ctx在這裡類似self，ctx的屬性可以在backward中呼叫。
    def forward(ctx, input, weight, bias=None)
 
:
        ctx.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


    @staticmethod
    def backward(ctx, grad_output): 
        input, weight, bias = ctx.saved_variables
        grad_input = grad_weight = grad_bias = None
 


        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0).squeeze(0)

        return grad_input, grad_weight, grad_bias
 

然後擴充套件module就很簡單，需要過載 nn.Module中的__init__和__forward__,

class Linear(nn.Module):
    def __init__(self, input_features, output_features, bias=True):
        super(Linear, self).__init__()
        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
        # 這個很重要！ Parameters是預設需要梯度的！
        # 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(output_features, input_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 LinearFunction.apply(input, self.weight, self.bias)

forward的說明

1. 雖然說一個網路的輸入是Variable形式，那麼每個網路層的輸出也是Variable形式。但是，當自定義autograd時，在forward中，所有的Variable引數將會轉成tensor！因此這裡的input也是tensor.在forward中可以任意操作。因為forward中是不需要Variable變數的，這是因為你自定義了backward方式。在傳入forward前，autograd engine會自動將Variable unpack成Tensor。
2. ctx是context，ctx.save_for_backward會將他們轉換為Variable形式。比如

    @staticmethod
    def backward(ctx, grad_output): 
        input, = ctx.saved_variables

此時input已經是需要grad的Variable了。
3. save_for_backward只能傳入Variable或是Tensor的變數，如果是其他型別的，可以用
ctx.xyz = xyz，使其在backward中可以用。

backward說明

grad_output是variable

其實這個沒啥好說的。就是預設情況下，你拿到grad_output時候，會發現它是一個Variable，至於requires_grad是否為True，取決於你在外面呼叫.backward或是.grad時候的那個Variable是不是需要grad的。如果那個Variable是需要grad的，那麼我們這裡反向的grad_ouput也是requires_grad為True，那麼我們甚至可以計算二階梯度。用WGAN-GP之類的。

backward中我能一開始就.data拿出資料進行操作嗎？

雖然自定義操作，但是原則上在backward中我們只能進行Variable的操作, 這顯然就要求我們在backward中的操作都是可自動求導的。因為預設情況下，
By default, the backward function should always work with Variable and create a proper graph (similarly to the forward function of an nn.Module).。所以如果我們的涉及到不可導的操作，那麼我們就不能在bacjward函式中建立一個正確的圖，值得注意的是，我們是

自動求導是根據每個op的backward建立的graph來進行的！

我去，知道真相的我眼淚掉下來！很奇怪吧！我的最初的想法就是forward記錄input如何被操作，然後backward就會自動反向，根據forward建立的圖進行！然而，當你print(type(input))時你竟然發現型別是Tensor，根本不是Variable！那怎麼記錄graph？然而真實情況竟然是自動求導竟然是在backward的操作中建立的圖！
這也就是為什麼我們需要在backward中用全部用variable來操作，而forward就沒必要，forward只需要用tensor操作就可以。

non-differential操作的backward怎麼寫？

當然你想個近似演算法，弄出來，一般直接backward中第一句就是grad_output = grad_output.data,這樣我們就無法進行建立正確的graph了。

#加一個這個
from torch.autograd.function import once_differentiable



@staticmethod
@once_differentiable
def backward(ctx, grad_output):
    print(type(grad_output))# 此時你會驚奇的發現，竟然是Tensor了！
    # 做點其他的操作得到grad_output_changed
    grad_input = grad_output_changed
    return grad_input

因為我們在backward中已經是直接拿出data進行操作的了，所以我們直接得到Tensor型別返回就行！！