[源码解析] PyTorch分布式优化器(1)

您所在的位置：网站首页 › pytorch源码解析setup与pybind › [源码解析] PyTorch分布式优化器(1)

[源码解析] PyTorch分布式优化器(1)

#[源码解析] PyTorch分布式优化器(1)| 来源: 网络整理| 查看: 265

0x00 摘要

我们接下来通过几篇文章来看看分布式优化器。本系列分为三篇文章，分别是基石篇，DP/DDP/Horovod 之中数据并行的优化器，PyTorch 分布式优化器，按照深度递进。

本文是基石篇，通过本文，大家可以了解到模型的构造，优化器的基本原理，两者之间的交互，如何优化更新模型等等，这为后面的逐级分析打下了一个基础。

PyTorch分布式其他文章如下：

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

[源码解析] PyTorch 如何使用GPU

[源码解析] 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) ---- 引擎(下)

为了更好的说明，本文代码会依据具体情况来进行相应精简。

0x01 从问题出发

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

1.1 示例

我们用个例子来看看如何进行训练。

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))) 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()

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

1.2 问题点

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

根据模型参数构建优化器

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

引擎计算梯度

如何保证 Linear 可以计算梯度？对于模型来说，计算出来的梯度怎么和 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 +----------------> | | | | | | | | | +------------------------------------------+ | | | 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 1.3 _parameters

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

1.3.1 构建

我们首先看看生成时候的特点：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) 1.3.2 归类

如果类的成员是从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 # 这里添加了 1.3.3 获取

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

比如：

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

输出：

torch.Size([10, 10]) torch.Size([10]) torch.Size([5, 10]) 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()) (20L,) (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 | | | forward / backward | | | Linear(10, 10)+--> ReLU +--> Linear(10, 5)| +----------------> | Compute gradient | | | | + | | para_iterator = parameters() | | | | | + ^ | | | | | | | | +------------------+ +-------------------------------------------+ | | | | gradient | | | 1 ??? | | 4 update v | | 2 ??? | | +----------------------------------------------------------------+ |SGD | | | | | | | | v | | | + | ^ +--------> self.parameters = para_iterator(ToyModel._parameters) ---------> | | | | | | | | | +----------------------------------------------------------------+ | | | None: factory_kwargs = {'device': device, 'dtype': dtype} super(Linear, self).__init__() self.in_features = in_features self.out_features = out_features self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs)) if bias: self.bias = Parameter(torch.empty(out_features, **factory_kwargs)) else: self.register_parameter('bias', None) self.reset_parameters() def reset_parameters(self) -> None: init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 init.uniform_(self.bias, -bound, bound) def forward(self, input: Tensor) -> Tensor: return F.linear(input, self.weight, self.bias) def extra_repr(self) -> str: return 'in_features={}, out_features={}, bias={}'.format( self.in_features, self.out_features, self.bias is not None ) 1.4.3 解释

从前面简略计算图我们可以知道，torch.nn.Linear 的反向计算是 AddmmBackward。

struct TORCH_API AddmmBackward : public TraceableFunction { using TraceableFunction::TraceableFunction; variable_list apply(variable_list&& grads) override; std::string name() const override { return "AddmmBackward"; } void release_variables() override { std::lock_guard lock(mutex_); mat2_.reset_data(); mat1_.reset_data(); } std::vector mat1_sizes; std::vector mat1_strides; SavedVariable mat2_; at::Scalar alpha; SavedVariable mat1_; std::vector mat2_sizes; std::vector mat2_strides; at::Scalar beta; };

我们从代码之中找到了 addmm 的定义，其注释说明这是个矩阵乘法操作。

def addmm(mat: Tensor, mat1: Tensor, mat2: Tensor, beta: float = 1., alpha: float = 1.) -> Tensor: r""" This function does exact same thing as :func:`torch.addmm` in the forward, except that it supports backward for sparse matrix :attr:`mat1`. :attr:`mat1` need to have `sparse_dim = 2`. Note that the gradients of :attr:`mat1` is a coalesced sparse tensor. Args: mat (Tensor): a dense matrix to be added mat1 (Tensor): a sparse matrix to be multiplied mat2 (Tensor): a dense matrix to be multiplied beta (Number, optional): multiplier for :attr:`mat` (:math:`\beta`) alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) """ return torch._sparse_addmm(mat, mat1, mat2, beta=beta, alpha=alpha)

目前我们可以继续拓展。

Linear 里面的 weight，bias 都是 Parameter 类型。 Parameter 构造函数中参数 requires_grad=True。这么设置就说明 Parameter 默认是需要计算梯度的。所以 Linear 的 weight，bias 就是需要引擎计算其梯度。 ToyModel 的 _parameters 成员变量通过 parameters 方法来获取，其返回的是一个Iterator。这个 iterator 作为参数用来构建 SGD 优化器。现在 SGD 优化器的 parameters 是一个指向 ToyModel._parameters 的 iterator。这说明优化器实际上是直接优化 ToyModel 的 _parameters，对于例子就是全连接层的参数，图上对应两个Linear 发出的指向 parameters() 的箭头。 +--------------------------------------------------+ +------------------+ | ToyModel | | Engine | | +-------------------+ +------------+ |forward / backward | | | | Linear(10, 10) +--> ReLU +-->+Linear(10,5)| +-----------------> | Compute gradient | | | | | | | | + | | | weight=Parameter | | weight | | | | | | | +----------+ | | | | | | | | bias=Parameter | | | bias | | +------------------+ | | | | | | | | | +-------------------+ | +--+---------+ | 2 | gradient | | | | | | | | | v | v v | ??? | para_iterator = parameters() | | + ^ | | | | | | | | | +--------------------------------------------------+ | | 1 ??? | | 4 update | | | | +----------------------------------------------------------------+ |SGD | | | | | | | | v | | | + | ^ +--------> self.parameters = para_iterator(ToyModel._parameters) +--------> | | | | | | | | | +----------------------------------------------------------------+ | | | ReLU +-->+Linear(10,5)| +-----------------> | Compute gradient | | | | | | | | + | | | weight=Parameter| | weight | | | | | | | +-----------+ | bias | | | | | | | bias=Parameter | | +--+---------+ | +------------------+ | | | | | | | | +------------------+ | | | 2 | gradient | v v | | | self._parameters | v | + | ??? | | | | | | | v | | para_iterator = parameters() | | + ^ | | | | | | | | | +-------------------------------------------------+ | | 1 ??? | | 4 update | | +----------------------------------------------------------------+ |SGD | | | | | | | | v | | | + | ^ +-------> self.param_groups = para_iterator(ToyModel._parameters) --------> | | | | | | | | | +----------------------------------------------------------------+ | | | > optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) >>> optimizer.zero_grad() >>> loss_fn(model(input), target).backward() >>> optimizer.step()

PyTorch SGD with Momentum/Nesterov 的实现与Sutskever et. al.和其他框架的实现不同。

比如 PyTorch 使用如下方法来实现 Momentum 的特殊例子：

vt+1=μ∗vt+gt+1,pt+1=pt−lr∗vt+1,\begin{aligned} v_{t+1} & = \mu * v_{t} + g_{t+1}, \\ p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, \end{aligned}vt+1pt+1=μ∗vt+gt+1,=pt−lr∗vt+1,

其他框架则使用：

vt+1=μ∗vt+lr∗gt+1,pt+1=pt−vt+1.\begin{aligned} v_{t+1} & = \mu * v_{t} + \text{lr} * g_{t+1}, \\ p_{t+1} & = p_{t} - v_{t+1}. \end{aligned}vt+1pt+1=μ∗vt+lr∗gt+1,=pt−vt+1. 3.3 step

step 方法的作用就是在一定的算法协助下，对变量进行优化。此方法主要完成一次模型参数的更新

@torch.no_grad() def step(self, closure=None): """Performs a single optimization step. Args: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ # 使用 closure 重新计算loss loss = None if closure is not None: with torch.enable_grad(): loss = closure() # 使用计算得到的梯度更新变量 # self.param_groups 就是我们传入的参数列表 for group in self.param_groups: # 每一个group是一个dict, 其包含每组参数所需的必要参数 params_with_grad = [] d_p_list = [] momentum_buffer_list = [] # 本组参数更新所必需的设置 weight_decay = group['weight_decay'] momentum = group['momentum'] dampening = group['dampening'] nesterov = group['nesterov'] lr = group['lr'] for p in group['params']: # 遍历本组所有需要更新的参数 if p.grad is not None: params_with_grad.append(p) d_p_list.append(p.grad) state = self.state[p] if 'momentum_buffer' not in state: momentum_buffer_list.append(None) else: momentum_buffer_list.append(state['momentum_buffer']) F.sgd(params_with_grad, d_p_list, momentum_buffer_list, weight_decay=weight_decay, momentum=momentum, lr=lr, dampening=dampening, nesterov=nesterov) # update momentum_buffers in state for p, momentum_buffer in zip(params_with_grad, momentum_buffer_list): state = self.state[p] state['momentum_buffer'] = momentum_buffer return loss

其中 sgd 函数如下：

def sgd(params: List[Tensor], d_p_list: List[Tensor], momentum_buffer_list: List[Optional[Tensor]], *, weight_decay: float, momentum: float, lr: float, dampening: float, nesterov: bool): r"""Functional API that performs SGD algorithm computation. See :class:`~torch.optim.SGD` for details. """ for i, param in enumerate(params): d_p = d_p_list[i] # 正则化及动量累积 if weight_decay != 0: d_p = d_p.add(param, alpha=weight_decay) if momentum != 0: buf = momentum_buffer_list[i] if buf is None: # 历史更新量 buf = torch.clone(d_p).detach() momentum_buffer_list[i] = buf else: # 通过buf更新了self.state buf.mul_(momentum).add_(d_p, alpha=1 - dampening) if nesterov: d_p = d_p.add(buf, alpha=momentum) else: d_p = buf # 更新当前组学习参数 w.data -= w.grad*lr param.add_(d_p, alpha=-lr) # add_ 会更改对象数值 3.4 变量解析

我们接下来对全局参数具体做以下解析。

3.4.1 lr

这就是学习率，大家熟知的概念。

3.4.2 dampening

dampening 作用到偏导数之上，用于动量SGD中调节当前梯度权重。

对应公式如下：

vt=vt−1∗momentum+gt∗(1−dampening)v_t = v_{t-1} * momentum + g_t * (1 - dampening)vt=vt−1∗momentum+gt∗(1−dampening)

对应代码则是：

buf.mul_(momentum).add_(d_p, alpha=1 - dampening) 3.4.3 weight_decay

weight_decay是 L2 penalty系数，用当前可学习参数p的值修改偏导数。

待更新的可学习参数p的偏导数就是

gt=gt+(p∗weight_decay)g_t = g_t + ( p * weight\_decay)gt=gt+(p∗weight_decay)

对应代码是：

if weight_decay != 0: d_p = d_p.add(param, alpha=weight_decay) 3.4.4 nesterov

是否启用nesterov动量，从pytorch源码来看，当nesterov为True时，在上述得到 v_t 的基础上又使用了一次momentum和v_t。

▽wJ(w)+m∗vt+1\bigtriangledown_{w}J(w) + m * v_{t+1}▽wJ(w)+m∗vt+1 if (nesterov) { d_p = d_p.add(buf, momentum); } else { d_p = buf; } 3.4.5 Momentum

Momentum ：来源于物理学，翻译为动量或则冲量。作用是把上次更新于当前梯度结合来进行当前权值优化更新。

引入原因是：训练网络的初始化权值可能因为不合适而导致在训练过程之中出现局部最小值，没有找到全局最优。

而引入动量可以在一定程度上解决此问题。动量模拟物体运动时候的惯性，表示力对时间的积累效应。更新时候在一定程度之上保持以前更新的方向，同时结合当前梯度来调整更新的方向。动量越大，转换为势能的能量越大，可以增加稳定性，也能更快的学习，从而越有可能摆脱局部凹区域，进入全局凹区域。

原生权重更新公式如下：

w=w−Lr∗dww = w - Lr * dww=w−Lr∗dw

这里 w 是权重，Lr 是学习率，dw 是 w 的导数。

引入momentum之后的权重更新公式如下：

v=momentum∗v−Lr∗dww=w+vv= momentum*v - Lr*dw \\w = w + vv=momentum∗v−Lr∗dww=w+v

这里 momentum 是动量，v 是速度。这个公式的意思就是加上上次更新的 v 与 momentum 的乘积。当本次梯度下降 -Lr * dw 的方向与上次更新 v 的方向相同，则上次更新 v 可以起到正向加速作用。当本次梯度下降 -Lr * dw 的方向与上次更新 v 的方向相反，则上次更新 v 可以起到减速作用。

代码对应如下：

if momentum != 0: buf = momentum_buffer_list[i] if buf is None: buf = torch.clone(d_p).detach() momentum_buffer_list[i] = buf else: buf.mul_(momentum).add_(d_p, alpha=1 - dampening) if nesterov: d_p = d_p.add(buf, alpha=momentum) else: d_p = buf 0x04 可视化 4.1 目前问题

到目前为止，我们还是有几个问题没有解决，就是下面下划线之处。

根据模型参数构建优化器

采用 optimizer = optim.SGD(params=net.parameters(), lr = 1) 进行构造，这样看起来 params 被赋值到优化器的内部成员变量之上（我们假定是叫parameters）。模型包括两个全连结层 Linear，这些层如何更新参数？？？ Linear 里面的 weight，bias 都是 Parameter 类型。 Parameter 构造函数中参数 requires_grad=True。这么设置就说明 Parameter 默认是需要计算梯度的。所以 Linear 的 weight，bias 就是需要引擎计算其梯度。 ToyModel 的 _parameters 成员变量通过 parameters 方法来获取，其返回的是一个Iterator。这个 iterator 作为参数用来构建 SGD 优化器。现在 SGD 优化器的 parameters 是一个指向 ToyModel._parameters 的 iterator。这说明优化器实际上是直接优化 ToyModel 的 _parameters。

引擎计算梯度

如何保证 Linear 可以计算梯度？ weight，bias 都是 Parameter 类型，默认是需要计算梯度的。 2) 对于模型来说，计算出来的梯度怎么和 Linear 参数对应起来？引擎计算出来的这些梯度累积在哪里？？？

优化器优化参数：

3) 调用 step 进行优化，优化目标是优化器内部成员变量 self.parameters。 self.parameters 是一个指向 ToyModel._parameters 的 iterator。这说明优化器实际上是直接优化 ToyModel 的 _parameters。

优化器更新模型：

4) 优化目标（self.parameters）的更新实际上就是直接作用到模型参数（比如 Linear）之上。

我们打印 outputs 看看，可以看到其 next_functions 实际是有三个，说明前面的图例是我们简化的，我们需要再做进一步可视化。

outputs = {Tensor: 10} T = {Tensor: 5} data = {Tensor: 10} device = {device} cpu dtype = {dtype} torch.float32 grad = {NoneType} None grad_fn = {AddmmBackward} metadata = {dict: 0} {} next_functions = {tuple: 3} 0 = {tuple: 2} (, 0) 1 = {tuple: 2} (, 0) 2 = {tuple: 2} (, 0) __len__ = {int} 3 requires_grad = {bool} True is_cuda = {bool} False is_leaf = {bool} False is_meta = {bool} False is_mkldnn = {bool} False is_mlc = {bool} False is_quantized = {bool} False is_sparse = {bool} False is_sparse_csr = {bool} False is_vulkan = {bool} False is_xpu = {bool} False layout = {layout} torch.strided name = {NoneType} None names = {tuple: 2} (None, None) ndim = {int} 2 output_nr = {int} 0 requires_grad = {bool} True 4.2 PyTorchViz可视化网络

我们采用PyTorchViz来展示网络。

先安装库：

pip install torchviz

然后添加代码可视化，我们使用可视化函数make_dot()来获取绘图对象。运行之后，代码相同根目录下的data文件夹里会生成一个.gv文件和一个.png文件，.gv文件是Graphviz工具生成图片的脚本代码，.png是.gv文件编译生成的图片。默认情况下程序会自动打开.png文件。

import torch import torch.nn as nn import torch.optim as optim from torchviz import make_dot 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))) net = ToyModel() print(net) # 顺便打印一下看看 optimizer = optim.SGD(params=net.parameters(), lr = 1) optimizer.zero_grad() input = torch.randn(10,10) outputs = net(input) outputs.backward(outputs) optimizer.step() NetVis = make_dot(outputs, params=dict(list(net.named_parameters()) + [('x', input)])) NetVis.format = "bmp" # 文件格式 NetVis.directory = "data" # 文件生成的文件夹 NetVis.view() # 生成文件

输出。

ToyModel( (net1): Linear(in_features=10, out_features=10, bias=True) (relu): ReLU() (net2): Linear(in_features=10, out_features=5, bias=True) )

图例如下：

我们发现，之前的简略图忽略了 AccumulateGrad 这个关键环节，我们接下来就分析一下。

0x05 AccumulateGrad 5.1 原理

我们首先来概述一下 PyTorch 相关原理知识。

从概念上讲，autograd 记录了一个计算图。图中节点分为两种：叶子节点和非叶子节点。

由用户创建的节点称为叶子节点，比如：

a=torch.tensor([1.0]) 运行时变量为： a = {Tensor: 1} tensor([1.]) T = {Tensor: 1} tensor([1.]) data = {Tensor: 1} tensor([1.]) device = {device} cpu dtype = {dtype} torch.float32 grad = {NoneType} None grad_fn = {NoneType} None is_cuda = {bool} False is_leaf = {bool} True requires_grad = {bool} False

但是此时 a 不能求导，在创建张量时，如果设置 requires_grad 为Ture，那么 Pytorch 才知道需要对该张量进行自动求导。

a=torch.tensor([1.0], requires_grad = True) 运行时变量为： a = {Tensor: 1} tensor([1.], requires_grad=True) T = {Tensor: 1} tensor([1.], grad_fn=) data = {Tensor: 1} tensor([1.]) device = {device} cpu dtype = {dtype} torch.float32 grad = {NoneType} None grad_fn = {NoneType} None is_cuda = {bool} False is_leaf = {bool} True requires_grad = {bool} True shape = {Size: 1} 1

PyTorch会记录对该张量的每一步操作历史，从而生成一个概念上的有向无环图，该无环图的叶子节点是模型的输入张量，其根为模型的输出张量。用户不需要对图的所有执行路径进行编码，因为用户运行的就是用户后来想微分的。通过从根到叶跟踪此图形，用户可以使用链式求导规则来自动计算梯度。

在内部实现上看，autograd 将此图表示为一个“Function” 或者说是"Node" 对象（真正的表达式）的图，该图可以使用apply方法来进行求值。

反向传播时候，autograd 引擎沿着从根节点（就是前向传播的输出节点）溯源这个图，这样就可以利用链式求导法则计算所有叶子节点的梯度。每一个前向传播操作函数都有一个反向传播函数与之对应，这个反向传播函数用来计算每个variable的梯度。

反向图之中，需要求导的叶子节点张量对应的反向传播计算函数就是AccumulateGrad，其梯度是累加的，多次求导都会在这个张量的导数上累积，比如：

a=torch.tensor([5.0], requires_grad = True) b = torch.tensor([3.0], requires_grad = True) c = a + b

对应的是：

对应我们的示例，Linear 实例都是用户显式定义的，所有都是叶子节点。

5.2 AccumulateGrad 5.2.1 定义

定义如下，accumulateGrad 实际就是：

先累积梯度。再调用传入的 update_grad 函数来更新梯度。 struct TORCH_API AccumulateGrad : public Node { explicit AccumulateGrad(Variable variable_); variable_list apply(variable_list&& grads) override; static at::Tensor callHooks( const Variable& variable, at::Tensor new_grad) { for (auto& hook : impl::hooks(variable)) { new_grad = (*hook)({new_grad})[0]; } return new_grad; } template static void accumulateGrad( const Variable& variable, at::Tensor& variable_grad, const at::Tensor& new_grad, size_t num_expected_refs, const T& update_grad) { // 传入的更新梯度函数 if (!variable_grad.defined()) { // 忽略 } else if (!GradMode::is_enabled()) { if (variable_grad.is_sparse() && !new_grad.is_sparse()) { auto result = new_grad + variable_grad; update_grad(std::move(result)); } else if (!at::inplaceIsVmapCompatible(variable_grad, new_grad)) { auto result = variable_grad + new_grad; update_grad(std::move(result)); } else { variable_grad += new_grad; // 进行累积 } } else { at::Tensor result; if (variable_grad.is_sparse() && !new_grad.is_sparse()) { // CPU backend throws an error on sparse + dense, so prefer dense + sparse here. result = new_grad + variable_grad; // 进行累积 } else { // Assumes operator+ result typically matches strides of first arg, // and hopes variable_grad was originally created obeying layout contract. result = variable_grad + new_grad; // 进行累积 } update_grad(std::move(result)); } } Variable variable; }; 5.2.2 apply

当调用 apply 时候，有两个注意点：

传入的更新函数就是 { grad = std::move(grad_update); } 更新梯度。 mutable_grad 得到的是张量的梯度成员变量。 Tensor& mutable_grad() const { return impl_->mutable_grad(); } /// Accesses the gradient `Variable` of this `Variable`. Variable& mutable_grad() override { return grad_; }

具体代码如下：

auto AccumulateGrad::apply(variable_list&& grads) -> variable_list { check_input_variables("AccumulateGrad", grads, 1, 0); if (!grads[0].defined()) return {}; if (variable.grad_fn()) throw std::logic_error( "leaf variable has been moved into the graph interior"); if (!variable.requires_grad()) return {}; at::Tensor new_grad = callHooks(variable, std::move(grads[0])); std::lock_guard lock(mutex_); at::Tensor& grad = variable.mutable_grad(); // 得到变量的mutable_grad accumulateGrad( variable, grad, new_grad, 1 + !post_hooks().empty() /* num_expected_refs */, [&grad](at::Tensor&& grad_update) { grad = std::move(grad_update); }); return variable_list(); }

具体流程图逻辑如下：

AccumulateGrad Tensor AutogradMeta + + + | | | | | | | | | v | | apply(update_grad) | | + | | | | | | | | | | | v | | accumulateGrad | | + | | | | | | result = variable_grad + new_grad | | | | | v result v v update_grad +----------------------------> mutable_grad +---> grad_

或者如下，对于一个叶子张量，反向计算时候会调用AccumulateGrad进行累积梯度，然后更新到叶子张量的 grad_ 之中：

+----------------------------------------------+ +-------------------------+ |Tensor | |TensorImpl | | | | | | | bridge | | | impl_ +-----------> | autograd_meta_ +---------+ | | | | | | | | | | +----------------------------------------------+ +-------------------------+ | | | | +-------------------------+ | | AutogradMeta | | | | | | | | | | apply(grads) { | | | | | | grad_accumulator_ | | accumulateGrad(new_grad) { | | | | | | | | result = variable_grad + new_grad | | | update | | | grad_ 引擎计算梯度 ---> 优化器优化参数 ---> 优化器更新模型这个顺序来总结。

根据模型参数构建优化器

采用 optimizer = optim.SGD(params=net.parameters(), lr = 1) 进行构造，这样 params 被赋值到优化器的内部成员变量 param_groups 之上。模型包括两个 Linear，这些层如何更新参数？ Linear 里面的 weight，bias 都是 Parameter 类型。 Parameter 构造函数中参数 requires_grad=True。这么设置就说明 Parameter 默认是需要计算梯度的。所以 Linear 的 weight，bias 就是需要引擎计算其梯度。 weight，bias 被添加到 ToyModel 的 _parameters 成员变量之中。 ToyModel 的 _parameters 成员变量通过 parameters 方法来获取，其返回的是一个Iterator。用这个 iterator 作为参数用来构建 SGD 优化器。现在 SGD 优化器的 parameters 是一个指向 ToyModel._parameters 的 iterator。这说明优化器实际上是直接优化 ToyModel 的 _parameters。所以优化器就是直接优化更新 Linear 的 weight 和 bias。其实优化器就是一套代码而已，具体优化哪些东西，需要在构建时候指定，优化一个模型的参数也行，优化用户自己指定的其他变量也行。

引擎计算梯度

如何保证 Linear 可以计算梯度？ weight，bias 都是 Parameter 类型，默认是需要计算梯度的。所以计算 weight，bias 梯度。对于模型来说，计算出来的梯度怎么和 Linear 参数对应起来？引擎计算出来的这些梯度累积在哪里？对应我们的示例，Linear 实例都是用户显式定义的，所以都是叶子节点。叶子节点通过 AccumulateGrad 把梯度累积在模型参数张量 autograd_meta_.grad_ 之中。

优化器优化参数：

4) 调用 step 进行优化，优化目标是优化器内部成员变量 self.parameters。 self.parameters 是一个指向 ToyModel._parameters 的 iterator。这说明优化器实际上是直接优化 ToyModel 的 _parameters。

优化器更新模型：

5) 优化目标（self.parameters）的更新实际上就是直接作用到模型参数（比如 Linear 的 weight，bias）之上。

具体如图：

+---------------------------------------------------------------------+ | ToyModel | | +---------------------------------+ +------------+ | +------------------+ | | Linear(10, 10) +------> ReLU +-->+Linear(10,5)| | | Engine | | | | | | |forward / backward | | | | weight=Parameter | | weight | +-----------------> | Compute gradient | | | +---------------+ | bias | | | + | | | +----------------------------+ | | +--+---------+ | | | | | | | bias=Parameter | | | | | | | | | | | | | | | | +------------------+ | | | | | | | | 3 accumulate | | | | autograd_meta_.grad_ | | | | | | | | | +---------------------------------------------------------------------------+ | | |

【本文地址】

[源码解析] PyTorch分布式优化器(1)

[源码解析] PyTorch分布式优化器(1)

今日新闻

推荐新闻