一、解析層的結構

首先我們通過分析官方的原始碼瞭解一下什麼是層，它包含哪些結構，成員是啥等。

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
        self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(output_features))
        else:
            self.register_parameter('bias', None)
        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):
        return LinearFunction.apply(input, self.weight, self.bias)
 

上面的原始碼是官方的線性層實現方式。

引數含義：

• input_features是輸入向量長度，output_features是輸出向量的長度
• input呼叫該類時的輸入

Linear層包含兩個內部引數，也就是我們說的層的權重，weight和bias。兩個函式建構函式__init__和前向傳播函式forward。

我們可以得到以下結論：

• pytorch的層繼承自nn.module類
• 層至少包含兩個函式成員__init__和前向傳播函式forward（如果自定義的操作不可導，還需要實現反向傳播的backward）
• 如果該層含有權重，那麼權重必須是nn.Parameter型別，關於Tensor和Variable（0.3版本之前）與Parameter的區別請參看之前
  部落格。簡單說就是Parameter預設需要求導，其他兩個型別則不會。
• 可能的話，為自己定義的新層提供預設的引數初始化，以防使用過程中忘記初始化操作。

二、示例

下面我們實現一個簡單的層，輸入[x,y],輸出為z，實現z=a*x+b*y的功能，並通過網路自動學習到引數a，b。

首先分析一下我們所要實現的功能z=a*x+b*y，其中有兩個要學習的引數，a和b。假設輸入為一個1*2向量，為了利用pytorch的乘法，我們將a，b合起來定義為[1,2]的向量，型別為Parmeter（所有代權重層中引數的型別）。為了更普適性完成 的任務，我們將引數的形狀設為在定義時指定。

層的定義

##################################################################
####in_features->該層的形狀，e.g 引數為a，b，則(1,2);為a,b,c，則(1,3)
####reset_parameters()權重預設初始化函式
####forward自己定義的操作
####input->呼叫該層時的輸入  shape->[n,1,2]
#################################################################


class weight_pool(nn.Module):
    def __init__(self, in_features):            
        super(weight_pool, self).__init__()
        self.in_features = in_features
        self.weight = nn.Parameter(torch.Tensor(self.in_features))
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(0))
        self.weight.data.uniform_(-stdv, stdv)
        
    def forward(self, input):
        x = input * self.weight
        x = x.sum(dim=1,keepdim=True)
        return x

此時就完成了我們定義的新層。下面通過幾組測試來檢驗我們的新層是否具有學習功能。

1.task1 學習z = x + y

由於需要學習的引數只有兩個，理論上只需要兩組資料就能完成學習。但是為了更普適，我們輸入了五組訓練資料

x,y = [1.0,2.0], [1.0,3.0], [2.0,3.0], [3.0,4.0], [9.0,10.0]
z = [3.0], [4.0], [5.0], [7.0], [19.0]

損失函式用MSELoss，學習率0.01，SGD方法，迭代10個epoch。

網路定義如下：

class MyNet(nn.Module):

    def __init__(self):
        super(MyNet, self).__init__()
        self.wpool = weight_pool((1,2))
    def forward(self, x):
        x = self.wpool(x)
        return x

下面展示前10個epoch的loss曲線圖，可以看到2個epoch時網路已經接近收斂。此時學習到的引數為1.0324和0.9731，非常接近我們最理想的引數1和1。

為了測試網路的極限效能，我們直接迭代1000個epoch，可以看到此時學習到的引數就是最理想的引數1和1.

task2  學習z = x + 3*y

同樣損失函式用MSELoss，學習率0.01，SGD方法，迭代10個epoch

此時由於目標函式較複雜，10個epoch之後，得到結果僅為1.7896和2.3437，與我們的理想結果1和3相差較遠。

同樣的，我們將迭代次數增加到1000個epoch，可以看到學習到的結果為1.0077和2.9936，與目標結果1和3基本一致。

總結：

通過上面的例子我們可以證明我們自己寫的層具有和pytorch原有的層一樣，具有學習能力，可以完成我們的目標。最後再總結一下新實現層的要點：

• 繼承nn.module類
• 層的權重型別為Parameter
• 至少實現兩個函式__init__和forward
• 自定義操作如不可導，需要實現backward函式。

完整的工程可在此下載。