pytorch源码解析：Python层 pytorchmodule源码

╰+哭是因爲堅強的太久メ 2023-06-05 03:24 17阅读 0赞

尝试使用了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][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组织成函数表达式（计算图）：

![Image 1][]

*  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及其参数以字典的方式保存。

![Image 2][]

### 前向传播、反向传播 ###

前向传播

例子中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

![1432127-20190830112952220-1979115827.jpg][]

转载于:https://www.cnblogs.com/jfdwd/p/11434087.html

[Deep Learning with PyTorch_ A 60 Minute Blitz]: http://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html
[Image 1]: /images/20230601/68ea60dcc93a45aaa88dd548270901b8.png
[Image 2]: /images/20230601/1751c52d4468499598a6b60db1ad6711.png
[1432127-20190830112952220-1979115827.jpg]: /images/20230601/b3db417f639f4a2a9349311154a98fb2.png