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 承載了實際的功能，定義了前向和後向的計算邏輯。
Module 是任何神經網路的基類，pytorch 中所有模型都必需是 Module 的子類。 Module 可以套嵌，構成樹狀結構。一個 Module 可以通過將其他 Module 做為屬性的方式，完成套嵌。
注意：真到目前（04/2018），pytorch 這部分的介面都沒穩定下來，下面的闡述已經和最新版本不一致，甚至不正確。在介面最終穩定之前，內容不再更新，請直接查閱 pytorch 的最新原始碼。
下面以最簡單的 MLP 網路結構為例，介紹下如何實現自定義網路結構。完整程式碼可以參見repo。
2.1 Function
注：為支援高階導數（i.e. 梯度的梯度），pytorch 0.2 收入新的定義 Function 的機制。如果不考慮高階，舊的方法依然 work。
Function 是 pytorch 自動求導機制的核心類。Function 是無引數或者說無狀態的，它只負責接收輸入，返回相應的輸出；對於反向，它接收輸出相應的梯度，返回輸入相應的梯度。
這裡我們只關注如何自定義 Function。Function 的定義見原始碼。下面是簡化的程式碼段：
class Function(object):
def forward(self, *input):
raise NotImplementedError
def backward(self, *grad_output):
raise NotImplementedErrorforward 和 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.saved_tensors
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為一個元素為 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），基本的計算公式如下：
圖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 output3.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 是一個非常優雅的深度學習框架。