Pytorch 學習筆記之自定義 Module

Pytorch 學習筆記之自定義 Module

pytorch 是一個基於 python 的深度學習庫。pytorch 源碼庫的抽象層次少，結構清晰，代碼量適中。相比於非常工程化的 tensorflow，pytorch 是一個更易入手的，非常棒的深度學習框架。

對於系統學習 pytorch，官方提供了非常好的入門教程 ，同時還提供了面向深度學習的示例，同時熱心網友分享了更簡潔的示例。

1. overview

不同於 theano，tensorflow 等低層程序庫，或者 keras、sonnet 等高層 wrapper，pytorch 是一種自成體系的深度學習庫（圖1）。

圖1. 幾種深度學習程序庫對比

如圖2所示，pytorch 由低層到上層主要有三大塊功能模塊。

圖2. pytorch 主要功能模塊

1.1 張量計算引擎（tensor computation）

Tensor 計算引擎，類似 numpy 和 matlab，基本對象是tensor（類比 numpy 中的 ndarray 或 matlab 中的 array）。除提供基於 CPU 的常用操作的實現外，pytorch 還提供了高效的 GPU 實現，這對於深度學習至關重要。

1.2 自動求導機制（autograd）

由於深度學習模型日趨復雜，因此，對自動求導的支持對於學習框架變得必不可少。pytorch 采用了動態求導機制，使用類似方法的框架包括： chainer，dynet。作為對比，theano，tensorflow 采用靜態自動求導機制。

1.3 神經網絡的高層庫（NN）

pytorch 還提供了高層的神經網絡模塊。對於常用的網絡結構，如全連接、卷積、RNN 等。同時，pytorch 還提供了常用的目標函數、optimizer 及參數初始化方法。

這里，我們重點關注如何自定義神經網絡結構。

2. 自定義 Module

圖3. pytorch Module

module 是 pytorch 組織神經網絡的基本方式。Module 包含了模型的參數以及計算邏輯。Function 承載了實際的功能，定義了前向和后向的計算邏輯。

下面以最簡單的 MLP 網絡結構為例，介紹下如何實現自定義網絡結構。完整代碼可以參見repo。

2.1 Function

Function 是 pytorch 自動求導機制的核心類。Function 是無參數或者說無狀態的，它只負責接收輸入，返回相應的輸出；對於反向，它接收輸出相應的梯度，返回輸入相應的梯度。

這里我們只關注如何自定義 Function。Function 的定義見源碼。下面是簡化的代碼段：

class Function(object): def forward(self, *input): raise NotImplementedError def backward(self, *grad_output): raise NotImplementedError 

forward 和 backward 的輸入和輸出都是 Tensor 對象。

Function 對象是 callable 的，即可以通過()的方式進行調用。其中調用的輸入和輸出都為 Variable 對象。下面的代碼示例了如何實現一個 ReLU 激活函數並進行調用：

import torch
from torch.autograd import Function

class ReLUF(Function)： def forward(self, input): self.save_for_backward(input) output = input.clamp(min=0) return output def backward(self, output_grad): input = self.to_save[0] input_grad = output_grad.clone() input_grad[input < 0] = 0 return input_grad ## Test if __name__ == "__main__": from torch.autograd import Variable torch.manual_seed(1111) a = torch.randn(2, 3) va = Variable(a, requires_grad=True) vb = ReLUF()(va) print va.data, vb.data vb.backward(torch.ones(va.size())) print vb.grad.data, va.grad.data 

如果 backward 中需要用到 forward 的輸入，需要在 forward 中顯式的保存需要的輸入。在上面的代碼中，forward 利用self.save_for_backward函數，將輸入暫時保存，並在 backward 中利用saved_tensors (python tuple 對象) 取出。

顯然，forward 的輸入應該和 backward 的輸入相對應；同時，forward 的輸出應該和 backward 的輸入相匹配。

由於 Function 可能需要暫存 input tensor，因此，建議不復用 Function 對象，以避免遇到內存提前釋放的問題。如示例代碼所示，forward的每次調用都重新生成一個 ReLUF 對象，而不能在初始化時生成在 forward 中反復調用。

2.2 Module

類似於 Function，Module 對象也是 callable 是，輸入和輸出也是 Variable。不同的是，Module 是[可以]有參數的。Module 包含兩個主要部分：參數及計算邏輯（Function 調用）。由於ReLU激活函數沒有參數，這里我們以最基本的全連接層為例來說明如何自定義Module。

全連接層的運算邏輯定義如下 Function:

import torch from torch.autograd import Function class LinearF(Function): def forward(self, input, weight, bias=None): self.save_for_backward(input, weight, bias) output = torch.mm(input, weight.t()) if bias is not None: output += bias.unsqueeze(0).expand_as(output) return output def backward(self, grad_output): input, weight, bias = self.saved_tensors grad_input = grad_weight = grad_bias = None if self.needs_input_grad[0]: grad_input = torch.mm(grad_output, weight) if self.needs_input_grad[1]: grad_weight = torch.mm(grad_output.t(), input) if bias is not None and self.needs_input_grad[2]: grad_bias = grad_output.sum(0).squeeze(0) if bias is not None: return grad_input, grad_weight, grad_bias else: return grad_input, grad_weight 

needs_input_grad 為一個元素為 bool 型的 tuple，長度與 forward 的參數數量相同，用來標識各個輸入是否輸入計算梯度；對於無需梯度的輸入，可以減少不必要的計算。

Function（此處為 LinearF） 定義了基本的計算邏輯，Module 只需要在初始化時為參數分配內存空間，並在計算時，將參數傳遞給相應的 Function 對象。代碼如下：

import torch import torch.nn as nn class Linear(nn.Module): def __init__(self, in_features, out_features, bias=True): super(Linear, self).__init__() self.in_features = in_features self.out_features = out_features self.weight = nn.Parameter(torch.Tensor(out_features, in_features)) if bias: self.bias = nn.Parameter(torch.Tensor(out_features)) else: self.register_parameter('bias', None) def forward(self, input): return LinearF()(input, self.weight, self.bias) 

需要注意的是，參數是內存空間由 tensor 對象維護，但 tensor 需要包裝為一個Parameter 對象。Parameter 是 Variable 的特殊子類，僅有是不同是 Parameter 默認requires_grad為 True。Varaible 是自動求導機制的核心類，此處暫不介紹，參見教程。

3. 自定義循環神經網絡（RNN）

我們嘗試自己定義一個更復雜的 Module ——RNN。這里，我們只定義最基礎的 vanilla RNN（圖4），基本的計算公式如下：

ht=relu(W⋅x+U⋅ht−1)

圖4. RNN【來源】

更復雜的 LSTM、GRU 或者其他變種的實現也非常類似。

3.1 定義 Cell

import torch from torch.nn import Module, Parameter class RNNCell(Module): def __init__(self, input_size, hidden_size): super(RNNCell, self).__init__() self.input_size = input_size self.hidden_size = hidden_size self.weight_ih = Parameter(torch.Tensor(hidden_size, input_size)) self.weight_hh = Parameter(torch.Tensor(hidden_size, hidden_size)) self.bias_ih = Parameter(torch.Tensor(hidden_size)) self.bias_hh = Parameter(torch.Tensor(hidden_size)) self.reset_parameters() def reset_parameters(self): stdv = 1.0 / math.sqrt(self.hidden_size) for weight in self.parameters(): weight.data.uniform_(-stdv, stdv) def forward(self, input, h): output = LinearF()(input, self.weight_ih, self.bias_ih) + LinearF()(h, self.weight_hh, self.bias_hh) output = ReLUF()(output) return output 

3.2 定義完整的 RNN

import torch from torch.nn import Module class RNN(Moudule): def __init__(self, input_size, hidden_size): super(RNN, self).__init__() self.input_size = input_size self.hidden_size = hidden_size sef.cell = RNNCell(input_size, hidden_size) def forward(self, inputs, initial_state): time_steps = inputs.size(1) state = initial_state outputs = [] for t in range(time_steps): state = self.cell(inputs[:, t, :], state) outputs.append(state) return outputs 

可運行的完整代碼見repo。

討論

pytorch 的 Module 結構是傳承自 torch，這一點也同樣被 keras （functional API）所借鑒。 在 caffe 等一些[早期的]深度學習框架中，network 是由於若干 layer ，經由不同的拓撲結構組成的。而在 (pyt)torch 中沒有 layer 和 network 是區分，一切都是 callable 的 Module。Module 的調用的輸入和輸出都是 tensor （由 Variable 封裝），用戶可以非常自然的構造任意有向無環的網絡結構（DAG）。

同時， pytorch 的 autograd 機制封裝的比較淺，可以比較容易的定制反傳或修改梯度。這對有些算法是非常重要。

總之，僅就自定義算法而言，pytorch 是一個非常優雅的深度學習框架。