pytorch源碼解析：Python層 pytorchmodule源碼

嘗試使用了pytorch，相比其他深度學習框架，pytorch顯得簡潔易懂。花時間讀了部分源碼，主要結合簡單例子帶着問題閱讀，不涉及源碼中C拓展庫的實現。

一個簡單例子

實現單層softmax二分類，輸入特征維度為4，輸出為2，經過softmax函數得出輸入的類別概率。代碼示意：定義網絡結構；使用SGD優化；迭代一次，隨機初始化三個樣例，每個樣例四維特征，target分別為1,0,1；前向傳播，使用交叉熵計算loss；反向傳播，最后由優化算法更新權重，完成一次迭代。

1.  
   import torch
2.  
   import torch.nn as nn
3.  
   import torch.nn.functional as F
4.  
    
5.  
   class Net(nn.Module):
6.  
    
7.  
   def __init__(self):
8.  
   super(Net, self).__init__()
9.  
   self.linear = nn.Linear( 4, 2)
10.  
     
11.  
    def forward(self, input):
12.  
    out = F.softmax(self.linear(input))
13.  
    return out
14.  
     
15.  
    net = Net()
16.  
    sgd = torch.optim.SGD(net.parameters(), lr= 0.001)
17.  
    for epoch in range(1):
18.  
    features = torch.autograd.Variable(torch.randn( 3, 4), requires_grad=True)
19.  
    target = torch.autograd.Variable(torch.LongTensor([ 1, 0, 1]))
20.  
    sgd.zero_grad()
21.  
     
22.  
    out = net(features)
23.  
    loss = F.cross_entropy(out, target)
24.  
    loss.backward()
25.  
    sgd.step()

從上面的例子，帶着下面的問題閱讀源碼：

• pytorch的主要概念：Tensor、autograd、Variable、Function、Parameter、Module（Layers）、Optimizer；
• 自定義Module如何組織網絡結構和網絡參數；
• 前向傳播、反向傳播實現流程
• 優化算法類如何實現，如何和自定義Module聯系並更新參數。

pytorch的主要概念

pytorch的主要概念官網有很人性化的教程Deep Learning with PyTorch: A 60 Minute Blitz， 這里簡單概括這些概念：

Tensor

類似numpy的ndarrays，強化了可進行GPU計算的特性，由C拓展模塊實現。如上面的torch.randn(3, 4) 返回一個3*4的Tensor。和numpy一樣，也有一系列的Operation，如

1.  
   x = torch.rand( 5, 3)
2.  
   y = torch.rand( 5, 3)
3.  
   print x + y
4.  
   print torch.add(x, y)
5.  
   print x.add_(y)

Varaiable與autograd

Variable封裝了Tensor，包括了幾乎所有的Tensor可以使用的Operation方法，主要使用在自動求導(autograd)，Variable類繼承_C._VariableBase，由C拓展類定義實現。
Variable是autograd的計算單元，Variable通過Function組織成函數表達式（計算圖）：

• data 為其封裝的tensor值
• grad 為其求導后的值
• creator 為創建該Variable的Function，實現中grad_fn屬性則指向該Function。
  如：
  1.  
     import torch
  2.  
     from torch.autograd import Variable
  3.  
     x = Variable(torch.ones( 2, 2), requires_grad=True)
  4.  
     y = x + 2
  5.  
     print y.grad_fn
  6.  
     print "before backward: ", x.grad
  7.  
     y.backward()
  8.  
     print "after backward: ", x.grad

  輸出結果：

  1.  
     < torch.autograd.function.AddConstantBackward object at 0x7faa6f3bdd68>
  2.  
     before backward: None
  3.  
     after backward: Variable containing:
  4.  
     1
  5.  
     [torch.FloatTensor of size 1x1]

  調用y的backward方法，則會對創建y的Function計算圖中所有requires_grad=True的Variable求導（這里的x）。例子中顯然dy/dx = 1。

Parameter

   Parameter 為Variable的一個子類，后面還會涉及，大概兩點區別：

• 作為Module參數會被自動加入到該Module的參數列表中；
• 不能被volatile， 默認require gradient。

Module

Module為所有神經網絡模塊的父類，如開始的例子，Net繼承該類，____init____中指定網絡結構中的模塊，並重寫forward方法實現前向傳播得到指定輸入的輸出值，以此進行后面loss的計算和反向傳播。

Optimizer

Optimizer是所有優化算法的父類（SGD、Adam、...），____init____中傳入網絡的parameters, 子類實現父類step方法，完成對parameters的更新。

自定義Module

該部分說明自定義的Module是如何組織定義在構造函數中的子Module,以及自定義的parameters的保存形式，eg:

1.  
   class Net(nn.Module):
2.  
   def __init__(self):
3.  
   super(Net, self).__init__()
4.  
   self.linear = nn.Linear(4, 2)
5.  
    
6.  
   def forward(self, input):
7.  
   out = F.softmax( self.linear(input))
8.  
   return out

首先看構造函數，Module的構造函數初始化了Module的基本屬性，這里關注_parameters和_modules，兩個屬性初始化為OrderedDict()，pytorch重寫的有序字典類型。_parameters保存網絡的所有參數，_modules保存當前Module的子Module。
module.py：

1.  
   class Module(object):
2.  
    
3.  
   def __init__(self):
4.  
   self._parameters = OrderedDict()
5.  
   self._modules = OrderedDict()
6.  
   ...

下面來看自定義Net類中self.linear = nn.Linear(4, 2)語句和_modules、_parameters如何產生聯系，或者self.linear及其參數如何被添加到_modules、_parameters字典中。答案在Module的____setattr____方法，該Python內建方法會在類的屬性被賦值時調用。
module.py:

1.  
   def __setattr__(self, name, value):
2.  
   def remove_from(*dicts):
3.  
   for d in dicts:
4.  
   if name in d:
5.  
   del d[name]
6.  
    
7.  
   params = self.__dict__.get( '_parameters')
8.  
   if isinstance(value, Parameter): # ----------- <1>
9.  
   if params is None:
10.  
    raise AttributeError(
11.  
    "cannot assign parameters before Module.__init__() call")
12.  
    remove_from(self.__dict__, self._buffers, self._modules)
13.  
    self.register_parameter(name, value)
14.  
    elif params is not None and name in params:
15.  
    if value is not None:
16.  
    raise TypeError("cannot assign '{}' as parameter '{}' "
17.  
    "(torch.nn.Parameter or None expected)"
18.  
    .format(torch.typename(value), name))
19.  
    self.register_parameter(name, value)
20.  
    else:
21.  
    modules = self.__dict__.get( '_modules')
22.  
    if isinstance(value, Module):# ----------- <2>
23.  
    if modules is None:
24.  
    raise AttributeError(
25.  
    "cannot assign module before Module.__init__() call")
26.  
    remove_from(self.__dict__, self._parameters, self._buffers)
27.  
    modules[name] = value
28.  
    elif modules is not None and name in modules:
29.  
    if value is not None:
30.  
    raise TypeError("cannot assign '{}' as child module '{}' "
31.  
    "(torch.nn.Module or None expected)"
32.  
    .format(torch.typename(value), name))
33.  
    modules[name] = value
34.  
    ......
35.  
     

調用self.linear = nn.Linear(4, 2)時，父類____setattr____被調用，參數name為“linear”， value為nn.Linear(4, 2)，內建的Linear類同樣是Module的子類。所以<2>中的判斷為真，接着modules[name] = value，該linear被加入_modules字典。
同樣自定義Net類的參數即為其子模塊Linear的參數，下面看Linear的實現：
linear.py:

1.  
   class Linear(Module):
2.  
    
3.  
   def __init__(self, in_features, out_features, bias=True):
4.  
   super(Linear, self).__init__()
5.  
   self.in_features = in_features
6.  
   self.out_features = out_features
7.  
   self.weight = Parameter(torch.Tensor(out_features, in_features))
8.  
   if bias:
9.  
   self.bias = Parameter(torch.Tensor(out_features))
10.  
    else:
11.  
    self.register_parameter('bias', None)
12.  
    self.reset_parameters()
13.  
     
14.  
    def reset_parameters(self):
15.  
    stdv = 1. / math.sqrt(self.weight.size(1))
16.  
    self.weight.data.uniform_(-stdv, stdv)
17.  
    if self.bias is not None:
18.  
    self.bias.data.uniform_(-stdv, stdv)
19.  
     
20.  
    def forward(self, input):
21.  
    return F.linear(input, self.weight, self.bias)

同樣繼承Module類，____init____中參數為輸入輸出維度，是否需要bias參數。在self.weight = Parameter(torch.Tensor(out_features, in_features))的初始化時，同樣會調用父類Module的____setattr____， name為“weight”，value為Parameter，此時<1>判斷為真，調用self.register_parameter(name, value)，該方法中對參數進行合法性校驗后放入self._parameters字典中。

Linear在reset_parameters方法對權重進行了初始化。

最終可以得出結論自定義的Module以樹的形式組織子Module，子Module及其參數以字典的方式保存。

前向傳播、反向傳播

前向傳播

例子中out = net(features)實現了網絡的前向傳播，該語句會調用Module類的forward方法，該方法被繼承父類的子類實現。net(features)使用對象作為函數調用，會調用Python內建的____call____方法，Module重寫了該方法。

module.py:

1.  
   def __call__( self, *input, **kwargs):
2.  
   for hook in self._forward_pre_hooks.values():
3.  
   hook( self, input)
4.  
   result = self.forward(*input, **kwargs)
5.  
   for hook in self._forward_hooks.values():
6.  
   hook_result = hook( self, input, result)
7.  
   if hook_result is not None:
8.  
   raise RuntimeError(
9.  
   "forward hooks should never return any values, but '{}'"
10.  
    "didn't return None".format(hook))
11.  
    if len(self._backward_hooks) > 0:
12.  
    var = result
13.  
    while not isinstance(var, Variable):
14.  
    var = var[0]
15.  
    grad_fn = var.grad_fn
16.  
    if grad_fn is not None:
17.  
    for hook in self._backward_hooks.values():
18.  
    wrapper = functools.partial(hook, self)
19.  
    functools.update_wrapper(wrapper, hook)
20.  
    grad_fn.register_hook(wrapper)
21.  
    return result

____call____方法中調用result = self.forward(*input, **kwargs)前后會查看有無hook函數需要調用（預處理和后處理）。
例子中Net的forward方法中out = F.softmax(self.linear(input))，同樣會調用self.linear的forward方法F.linear(input, self.weight, self.bias)進行矩陣運算（仿射變換）。
functional.py:

1.  
   def linear(input, weight, bias=None):
2.  
   if input.dim() == 2 and bias is not None:
3.  
   # fused op is marginally faster
4.  
   return torch.addmm(bias, input, weight.t())
5.  
    
6.  
   output = input.matmul(weight.t())
7.  
   if bias is not None:
8.  
   output += bias
9.  
   return output

最終經過F.softmax，得到前向輸出結果。F.softmax和F.linear類似前面說到的Function（Parameters的表達式或計算圖）。

反向傳播

得到前向傳播結果后，計算loss = F.cross_entropy(out, target)，接下來反向傳播求導數d(loss)/d(weight)和d(loss)/d(bias)：

loss.backward() 

backward()方法同樣底層由C拓展，這里暫不深入，調用該方法后，loss計算圖中的所有Variable(這里linear的weight和bias)的grad被求出。

Optimizer參數更新

在計算出參數的grad后，需要根據優化算法對參數進行更新，不同的優化算法有不同的更新策略。
optimizer.py:

1.  
   class Optimizer(object):
2.  
    
3.  
   def __init__(self, params, defaults):
4.  
   if isinstance(params, Variable) or torch.is_tensor(params):
5.  
   raise TypeError("params argument given to the optimizer should be "
6.  
   "an iterable of Variables or dicts, but got " +
7.  
   torch.typename(params))
8.  
    
9.  
   self.state = defaultdict(dict)
10.  
    self.param_groups = list(params)
11.  
    ......
12.  
     
13.  
    def zero_grad(self):
14.  
    """Clears the gradients of all optimized :class:`Variable` s."""
15.  
    for group in self.param_groups:
16.  
    for p in group['params']:
17.  
    if p.grad is not None:
18.  
    if p.grad.volatile:
19.  
    p.grad.data.zero_()
20.  
    else:
21.  
    data = p.grad.data
22.  
    p.grad = Variable(data.new().resize_as_(data).zero_())
23.  
     
24.  
    def step(self, closure):
25.  
    """Performs a single optimization step (parameter update).
26.  
     
27.  
    Arguments:
28.  
    closure (callable): A closure that reevaluates the model and
29.  
    returns the loss. Optional for most optimizers.
30.  
    """
31.  
    raise NotImplementedError

Optimizer在init中將傳入的params保存到self.param_groups，另外兩個重要的方法zero_grad負責將參數的grad置零方便下次計算，step負責參數的更新，由子類實現。
以列子中的sgd = torch.optim.SGD(net.parameters(), lr=0.001)為例，其中net.parameters()返回Net參數的迭代器，為待優化參數；lr指定學習率。
SGD.py:

1.  
   class SGD(Optimizer):
2.  
    
3.  
   def __init__(self, params, lr=required, momentum=0, dampening=0,
4.  
   weight_decay=0, nesterov=False):
5.  
   defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
6.  
   weight_decay=weight_decay, nesterov=nesterov)
7.  
   if nesterov and (momentum <= 0 or dampening != 0):
8.  
   raise ValueError("Nesterov momentum requires a momentum and zero dampening")
9.  
   super(SGD, self).__init__(params, defaults)
10.  
     
11.  
    def __setstate__(self, state):
12.  
    super(SGD, self).__setstate__(state)
13.  
    for group in self.param_groups:
14.  
    group.setdefault( 'nesterov', False)
15.  
     
16.  
    def step(self, closure=None):
17.  
    """Performs a single optimization step.
18.  
     
19.  
    Arguments:
20.  
    closure (callable, optional): A closure that reevaluates the model
21.  
    and returns the loss.
22.  
    """
23.  
    loss = None
24.  
    if closure is not None:
25.  
    loss = closure()
26.  
     
27.  
    for group in self.param_groups:
28.  
    weight_decay = group[ 'weight_decay']
29.  
    momentum = group[ 'momentum']
30.  
    dampening = group[ 'dampening']
31.  
    nesterov = group[ 'nesterov']
32.  
     
33.  
    for p in group['params']:
34.  
    if p.grad is None:
35.  
    continue
36.  
    d_p = p.grad.data
37.  
    if weight_decay != 0:
38.  
    d_p.add_(weight_decay, p.data)
39.  
    if momentum != 0:
40.  
    param_state = self.state[p]
41.  
    if 'momentum_buffer' not in param_state:
42.  
    buf = param_state[ 'momentum_buffer'] = d_p.clone()
43.  
    else:
44.  
    buf = param_state[ 'momentum_buffer']
45.  
    buf.mul_(momentum).add_( 1 - dampening, d_p)
46.  
    if nesterov:
47.  
    d_p = d_p.add(momentum, buf)
48.  
    else:
49.  
    d_p = buf
50.  
     
51.  
    p.data.add_(-group[ 'lr'], d_p)
52.  
     
53.  
    return loss
54.  
     

SGD的step方法中，判斷是否使用權重衰減和動量更新，如果不使用，直接更新權重param := param - lr * d(param)。例子中調用sgd.step()后完成一次epoch。這里由於傳遞到Optimizer的參數集是可更改（mutable）的，step中對參數的更新同樣是Net中參數的更新。

小結

到此，根據一個簡單例子閱讀了pytorch中Python實現的部分源碼，沒有深入到底層Tensor、autograd等部分的C拓展實現，后面再繼續讀一讀C拓展部分的代碼。

轉自鏈接：https://www.jianshu.com/p/f5eb8c2e671c