让我们通过几篇文章来看看分布式优化器。本系列分为三篇文章，即基石文章，DP/DDP/Horovod 并行数据优化器，PyTorch 按深度递进分布式优化器。

本文是基石。通过本文，我们可以了解模型的结构、优化器的基本原理、两者之间的互动、如何优化和更新模型等，为以后的逐步分析奠定了基础。

PyTorch其他分布式文章如下：

自动微分(1)深度学习利器

自动微分(2)深度学习利器

深度学习利器的自动微分(3) — 示例解读

[源码解析]PyTorch如何实现前向传播(1) — 基础类(上)

[源码解析]PyTorch如何实现前向传播(2) — 基础类(下)

[源码解析] PyTorch如何实现前向传播(3) — 具体实现

[源码解析] Pytorch 如何实现后传播？ (1)---- 调用引擎

[源码解析] Pytorch 如何实现后传播？ (2)---- 发动机静态结构

[源码解析] Pytorch 如何实现后传播？ (3)---- 引擎动态逻辑

[源码解析] PyTorch 如何实现后传播？ (4)---- 具体算法

[源码解析] PyTorch 分布式(1)-历史和概述

[源码解析] PyTorch 分布式(2) ----- DataParallel(上)

[源码解析] PyTorch 分布式(3) ----- DataParallel(下)

[源码解析] PyTorch 分布式(4)-分布式应用基本概念

[源码解析] PyTorch分布式(5) ------ DistributedDataParallel 总述&如何使用

[源码解析] PyTorch分布式(6) —DistributedDataParallel – 初始化&store

[源码解析] PyTorch 分布式(7) ----- DistributedDataParallel 之进程组

[源码解析] PyTorch 分布式(8) -------- DistributedDataParallel之论文篇

[源码解析] PyTorch 分布式(9) ----- DistributedDataParallel 之初始化

[源码解析] PyTorch 分布式(10)------DistributedDataParallel 之 Reducer静态架构

[源码解析] PyTorch 分布式(11) ----- DistributedDataParallel 之 构建Reducer和Join操作

[源码解析] PyTorch 分布式(12) ----- DistributedDataParallel 之 前向传播

[源码解析] PyTorch 分布式(13) ----- DistributedDataParallel 之 反向传播

[源码解析] PyTorch 分布式 Autograd (1) ---- 设计

[源码解析] PyTorch 分布式 Autograd (2) ---- RPC基础

[源码解析] PyTorch 分布式 Autograd (3) ---- 上下文相关

[源码解析] PyTorch 分布式 Autograd (4) ---- 如何切入引擎？

[源码解析] PyTorch 分布式 Autograd (5) ---- 引擎(上)

[源码解析] PyTorch 分布式 Autograd (6) ---- 引擎(下)

为了更好地解释，本文的代码将根据具体情况进行简化。

快手八卦论文下图列出了原生训练过程和DDP/Horovod的对比，上面的 vanilla 是原生训练过程，其中 U 部分对应于优化器过程。传统优化器的主要功能是根据梯度进行优化&更新模型当前参数 ： w.data -= w.grad * lr。

让我们举个例子来看看如何训练。

class ToyModel(nn.Module):     def __init__(sel):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

net = ToyModel()
optimizer = optim.SGD(params=net.parameters(), lr = 1)
optimizer.zero_grad()
input = torch.randn(10,10)
outputs = net(input)
outputs.backward(outputs)
optimizer.step()

给出一个粗略的反向计算图如下 。

因为已经有了之前分析引擎等其他经历，所以我们结合之前得到的知识先整理出几个问题点，用来引导我们分析，我们按照 ：根据模型参数构建优化器 —> 引擎计算梯度 —> 优化器优化参数 —> 优化器更新模型 这个顺序来分析。我们知道是autograd引擎计算了梯度，这样问题就来了：

• 根据模型参数构建优化器

  • 采用 optimizer = optim.SGD(params=net.parameters(), lr = 1) 进行构造，这样看起来 params 被赋值到优化器的内部成员变量之上（我们假定是叫parameters）。
  • 1. 模型包括两个 Linear，这些层如何更新参数？

• 引擎计算梯度

  • 如何保证 Linear 可以计算梯度？
  • 2. 对于模型来说，计算出来的梯度怎么和 Linear 参数对应起来？引擎计算出来的这些梯度累积在哪里？

• 优化器优化参数：

  • 3. 调用 step 进行优化，优化目标是优化器内部成员变量 self.parameters。

• 优化器更新模型：

  • 4. 如何把优化目标（self.parameters）的更新反应到模型参数（比如 Linear）的更新上？

下面图之中的数字和问号就对应了上面4个问题。

      +-------------------------------------------+                    +------------------+
      |ToyModel                                   |                    | Engine           |
      |                                           | forward / backward |                  |
      | Linear(10, 10)+--> ReLU +--> Linear(10, 5)| +----------------> | Compute gradient |
      |                                           |                    |        +         |
      +-------------------+-----------------------+                    |        |         |
                          |                                            |        |         |
                    1 ??? | parameters()                               +------------------+
                          |                                                     |
                          |                                                     | gradient
                          |   ^                                                 |
                          |   |                                                 v
                          |   | 4 ???                                        2 ???
                          |   |
      +------------------------------------------+
      |SGD                |   |                  |
      |                   |   |                  |
      |                   v   +                  |
      |                                          |
^ +---------------> self.parameters  +---------------->
|     |                                          |    |
|     |                                          |    |
|     +------------------------------------------+    |
|                                                     |
<---------------------------------------------------+ v
                     3 step()

我们需要一步一步来分析。

因为优化器是优化更新模型的参数，所以我们首先介绍下模型相关信息。

在PyTorch如果定义一个模型，一般需要继承 nn.Module。

import torch
import torch.nn as nn
import torch.nn.functional as F

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

Module 定义如下：

class Module:
    r"""Base class for all neural network modules. Your models should also subclass this class. Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:: import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self): super(Model, self).__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:`to`, etc. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool """

    dump_patches: bool = False
    _version: int = 1
    training: bool
    _is_full_backward_hook: Optional[bool]

    def __init__(self):
        """ Initializes internal Module state, shared by both nn.Module and ScriptModule. """
        torch._C._log_api_usage_once("python.nn_module")

        self.training = True
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._non_persistent_buffers_set = set()
        self._backward_hooks = OrderedDict()
        self._is_full_backward_hook = None
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()

Module 内部有如下重要变量，大致可以分为如下三类。

基础类型：

• _parameters ：类型为张量的权重参数，用于前向和后向传播，保存模型就是保存这些参数。使用 parameters() 函数可以递归获取到模型所有参数，但是需要注意，parameters() 函数返回的是 iterator。
• _buffers : 存储一些需要持久化的非网络参数的变量，比如BN 的 running_mean。
• _modules : 存储类型为 Module 的变量，当后去一个模型的parameters 时候，PyTorch 通过递归遍历所有_modules来实现。

计算相关类型：

在模型计算时候，是按照如下顺序完成：

 _backward_hooks  ----> forward ----> _forward_hooks ----> _backward_hooks

具体如下：

• _forward_pre_hooks ：在 forward 之前运行，不会更改 forward 输入参数。

• _forward_hooks ：在 forward 之后运行，不会改变 forward 的输入和输出。

• _backward_hooks ：在 backward 之后运行，不会改变 backward 的输入和输出。

保存/加载相关：

以下是保存相关的，PyTorch 使用如下来保存 torch.save(cn.state_dict()…) ，使用 load_state_dict(state_dict) 来加载。

• _load_state_dict_pre_hooks : 在调用 _load_from_state_dict 加载模型时希望执行的操作。
• _state_dict_hooks ：在调用state_dict方法时希望执行的操作。

具体运行时候如下：

net = { 
        ToyModel} 
 T_destination = { 
        TypeVar} ~T_destination
 dump_patches = { 
        bool} False
 net1 = { 
        Linear} Linear(in_features=10, out_features=10, bias=True)
 net2 = { 
        Linear} Linear(in_features=10, out_features=5, bias=True)
 relu = { 
        ReLU} ReLU()
 training = { 
        bool} True
  _backward_hooks = { 
        OrderedDict: 0} OrderedDict()
  _buffers = { 
        OrderedDict: 0} OrderedDict()
  _forward_hooks = { 
        OrderedDict: 0} OrderedDict()
  _forward_pre_hooks = { 
        OrderedDict: 0} OrderedDict()
  _is_full_backward_hook = { 
        NoneType} None
  _load_state_dict_pre_hooks = { 
        OrderedDict: 0} OrderedDict()
  _modules = { 
        OrderedDict: 3} OrderedDict([('net1', Linear(in_features=10, out_features=10, bias=True)), ('relu', ReLU()), ('net2', Linear(in_features=10, out_features=5, bias=True))])
  _non_persistent_buffers_set = { 
        set: 0} set()
  _parameters = { 
        OrderedDict: 0} OrderedDict()
  _state_dict_hooks = { 
        OrderedDict: 0} OrderedDict()
  _version = { 
        int} 1

优化器是优化 _parameters，所以我们需要特殊了解一下。

我们首先看看生成时候的特点：requires_grad=True。参数这么设置，就说明 Parameter 就是需要计算梯度的。

因为张量默认是不需要求导的，requires_grad属性默认为False，如果某个节点 requires_grad 属性被设置为True，就说明其需要求导，并且所有依赖于它的节点 requires_grad 都为True。

class Parameter(torch.Tensor):
    r"""A kind of Tensor that is to be considered a module parameter. Parameters are :class:`~torch.Tensor` subclasses, that have a very special property when used with :class:`Module` s - when they're assigned as Module attributes they are automatically added to the list of its parameters, and will appear e.g. in :meth:`~Module.parameters` iterator. Assigning a Tensor doesn't have such effect. This is because one might want to cache some temporary state, like last hidden state of the RNN, in the model. If there was no such class as :class:`Parameter`, these temporaries would get registered too. Args: data (Tensor): parameter tensor. requires_grad (bool, optional): if the parameter requires gradient. See :ref:`locally-disable-grad-doc` for more details. Default: `True` """
    def __new__(cls, data=None, requires_grad=True): # 需要计算梯度
        if data is None:
            data = torch.tensor([])
        return torch.Tensor._make_subclass(cls, data, requires_grad)

如果类的成员是从Parameter类派生，那么nn.Module使用__setattr__机制把他们归属到_parameters 之中。比如Linear的weight和bias。

def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -> None:
    
    # 省略 .....
    
    params = self.__dict__.get('_parameters')
    if isinstance(value, Parameter):
        remove_from(self.__dict__, self._buffers, self._modules, self._non_persistent_buffers_set)
        self.register_parameter(name, value) # 
        

    def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
        r"""Adds a parameter to the module. The parameter can be accessed as an attribute using given name. Args: name (string): name of the parameter. The parameter can be accessed from this module using the given name param (Parameter): parameter to be added to the module. """
        
        # 省略各种校验

        if param is None:
            self._parameters[name] = None
        elif not isinstance(param, Parameter):
            raise TypeError("cannot assign '{}' object to parameter '{}' "
                            "(torch.nn.Parameter or None required)"
                            .format(torch.typename(param), name))
        elif param.grad_fn:
            raise ValueError(
                "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
                "parameters must be created explicitly. To express '{0}' "
                "as a function of another Tensor, compute the value in "
                "the forward() method.".format(name))
        else:
            self._parameters[name] = param # 这里添加了
        

我们无法直接获取到 _parameters 这个变量，只能通过 parameters 方法来获取，其返回的是一个Iterator。

比如：

for param in net.parameters():
    print(type(param), param.size())

输出：

<class 'torch.nn.parameter.Parameter'> torch.Size([10, 10])
<class 'torch.nn.parameter.Parameter'> torch.Size([10])
<class 'torch.nn.parameter.Parameter'> torch.Size([5, 10])
<class 'torch.nn.parameter.Parameter'> torch.Size([5])

parameters 代码如下。

def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters. This is typically passed to an optimizer. Args: recurse (bool): if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module. Yields: Parameter: module parameter Example:: >>> for param in model.parameters(): >>> print(type(param), param.size()) <class 'torch.Tensor'> (20L,) <class 'torch.Tensor'> (20L, 1L, 5L, 5L) """
    for name, param in self.named_parameters(recurse=recurse):
        yield param

再来看看 named_parameters，其核心是 module._parameters.items()，以列表返回可遍历的元组数组。

def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself. Args: prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module. Yields: (string, Parameter): Tuple containing the name and parameter Example:: >>> for name, param in self.named_parameters(): >>> if name in ['bias']: >>> print(param.size()) """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem    

需要注意，我们目前已经有了两个关键知识：

• Parameter 构造函数中参数 requires_grad=True。这么设置就说明 Parameter 默认就是需要计算梯度的。
• 通过 parameters 方法来获取，其返回的是一个Iterator。

所以之前图可以拓展一下，现在 SGD 的 parameters 是一个指向 ToyModel._parameters 的 iterator，这说明优化器实际上是直接优化 ToyModel 的 _parameters。所以我们可以去掉原来图之中 4) 对应的问号。

      +-------------------------------------------+                    +------------------+
      |ToyModel                                   |                    | Engine