基于Pytorch源码对SGD、momentum、Nesterov学习

目前神经网络的监督学习过程通常为：

1. 数据加载（load）进神经网络
2. 经过网络参数对数据的计算，得出预测值（predict）
3. 根据预测值与标注值（label）之间的差距，产生损失（loss）
4. 通过反向传播（BP：Back Propagation）对神经网络的各个参数产生梯度（gradient）
5. 依据特定的梯度下降算法（如SGD：Stochastic Gradient Descent随机梯度下降），基于梯度对参数进行更新
6. 前五步重复，直到网络参数收敛（convergence）

本文主要基于Pytorch的sgd源代码，对第五步的梯度下降算法进行讨论，并加深印象。

Pytorch中对优化器常见的使用方式：

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for i in range(epoch):
	...
	optimizer.zero_grad()
	loss = ...
	loss.backward()
	optimizer.step()

optimizer在每个epoch中的步骤有二。1、梯度置零；2、梯度更新。

---

zero_grad
zero_grad()函数是SGD类继承父类Optimizer的：

class SGD(Optimizer):
	...

父类Optimizer的zero_grad()中：1、遍历了网络的各个参数；2、对每个参数的梯度重置为0。
因为Pytorch中的梯度是以累加形式计算，并非直接覆盖上一次梯度。
因此每个epoch在backward前，需要显式地置零。

def zero_grad(self):
    r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
    for group in self.param_groups:
        for p in group['params']:
            if p.grad is not None:
                p.grad.detach_()
                p.grad.zero_()

---

step
SGD
梯度下降算法重点在于optimizer.step()：
当初始化optimizer时，只是简单地传入了parameters和lr（learning rate）

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

step的代码为：

for p in group['params']:
	if p.grad is None:
		continue
	d_p = p.grad
	# p = p - d_p * lr
    p.add_(d_p, alpha=-group['lr'])

这个是最基本的梯度下降算法。遍历每个参数，并减去它的梯度（grad）和学习率（lr）的积。
$$p_i=p_i-d\_p_i\ast lr$$

---

SGD+momentum
当初始化optimizer时，传入了momentum：

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

step的代码为：

for p in group['params']:
	if p.grad is None:
		continue
	d_p = p.grad

    if momentum != 0:
		param_state = self.state[p]
			# 如果没有momentum_buffer，就初始化一个
            if 'momentum_buffer' not in param_state:
            	# buf = d_p
            	buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
            else:
            	# buf = buf * momentum
            	# buf = buf + d_p
                buf = param_state['momentum_buffer']
                buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
            d_p = buf
            
	# p = p - d_p * lr = p - (momentum * buf + d_p) * lr
    p.add_(d_p, alpha=-group['lr'])

加上动量momentum后，梯度的更新会考虑过往的动量（momentum_buffer）。
结合过往的动量及本次的梯度，按照momentum初始化的比例（例如0.9），得出本轮要更新的幅度。
其中dampening也是超参数，取0到1，表示考虑多少本次的梯度。
$$\begin{gathered} buf=buf\ast momentum \\ buf=buf+d\_p \\ d\_p=buf \\ p=p-d\_p\ast lr \end{gathered}$$
也就对应于原公式的：
$$\begin{gathered} v_t=\gamma v_{t-1}+\eta {▽}_{\theta }J(\theta ) \\ \theta =\theta -v_t \end{gathered}$$
γ对应momentum，η对应（1 - dampening）。

---

Nesterov Accelerated Gradient（NAG）
SGD为最基本的参数更新。
momentum考虑了过往的累计动量，缓解了因梯度差异太大导致的抖动现象。
但是目前只考虑到了历史的情况，在参数即将更新到最优值时，却有可能因累计的动量过大，导致参数越过了最优状态。
如果参数在更新时，能具备一些先验的知识，在得知快到最优解时能适当降低更新幅度，适应性会更好。

先附上NAG的更新公式：
$$\begin{gathered} v_t=\gamma v_{t-1}+\eta {▽}_{\theta }J(\theta -\gamma v_{t-1}) \\ \theta =\theta -v_t \end{gathered}$$
公式中$\theta -\gamma v_{t-1}$近似当做更新后的参数的梯度，作为下一次梯度大小的先验知识。
初始化方式为：

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9, nesterov=True)

需要注意__init__处有这样的判断：

if nesterov and (momentum <= 0 or dampening != 0):
	raise ValueError("Nesterov momentum requires a momentum and zero dampening")

所以使用NAG时，momentum大于0，且dampening为0。

step为：

for p in group['params']:
	if p.grad is None:
		continue
	d_p = p.grad

    if momentum != 0:
		param_state = self.state[p]
			# 如果没有momentum_buffer，就初始化一个
            if 'momentum_buffer' not in param_state:
            	# buf = d_p
            	buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
            else:
            	# buf = buf * momentum + d_p
                buf = param_state['momentum_buffer']
                buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
            if nesterov:
            	# d_p = d_p + buf * momentum
                d_p = d_p.add(buf, alpha=momentum)
            else:
                d_p = buf
            
	# p = p - d_p * lr = p - (d_p + (buf * momentum + d_p) * momentum) * lr
    p.add_(d_p, alpha=-group['lr'])

是否加上Nesterov，源码上的差别只在于d_p的最终取值，即公式中的$v_t$。
从原本的
$$\begin{gathered} buf=\gamma t_{t-1}+▽J \\ {v}_t=buf \end{gathered}$$
变成
$$\begin{gathered} buf=\gamma t_{t-1}+▽J \\ {v}_t=\gamma buf+▽J \end{gathered}$$
说实话没太理解这部分是如何计算先验梯度，有可能是数学层面的trick？还是说有其他被忽略的地方。
待补坑，求解答~~
补坑
Pytorch源码中对于NAG的公式如下：
$$\begin{array}{rl}p& =p-lr\ast d\_p\\ & =p-(d\_p+(buf\ast momentum+d\_p)\ast momentum)\ast lr\\ & =p-(g+(v\ast m+g)\ast m)\ast lr\\ & =p-lr\ast g-lr\ast m\ast (m\ast v+g)\\ & =p+m\ast v-m\ast v-lr\ast g-lr\ast m\ast (m\ast v+g)\\ & =p+m\ast v-(m\ast v+lr\ast g)-lr\ast m\ast (m\ast v+g)\end{array}$$
再搬上Hinton大神的例子：
原本NAG起点为0，每一个step为：0→1→0→2，下一个step的起点为2，步骤为：2→3→2→右下角的点。
但实际上我们这种回溯是很没必要的，何不索性简化为0→1→2呢。
再往下思考，NAG的主要思想为提前走一步，再用这一步的梯度来更新参数。当前step的“提前走一步”其实就是上一step的“走多一步”，那不如每个step都走多一步，那下一个step就不用先试探性走，再更新参数了。于是NAG的步骤实现变为1→2→3。
每个step的起点变为1，1→2为参数更新，2→3可以理解为为下个step做准备。

再看回公式的最终结果
$p=p+m\ast v-(m\ast v+lr\ast g)-lr\ast m\ast (m\ast v+g)$
其中，$p+m\ast v$实际上是1→0，
$p+m\ast v-(m\ast v+lr\ast g)$是1→0→2，
$p+m\ast v-(m\ast v+lr\ast g)-lr\ast m\ast (m\ast v+g)$是1→0→2→3。
因为中间加上了 $+m\ast v-m\ast v$ ，实际上等价于源码中的1→2→3。

也算上数学上的trick了，优化了运算的步骤，缺没有丢失算法的核心思想。
坑补自知乎问题中，王占宇前辈的评论：怎么理解Pytorch中对Nesterov的实现？

---

注意
另外在源码的注释中有提到，Pytorch中的SGD实现，是以这种方式实现的（lr作为$v_{t+1}$整体的参数放在第二步）：
$$\begin{array}{rl}{v}_{t+1}& =\mu \ast v_t+g_{t+1},\\ p_{t+1}& =p_t-\text{lr}\ast v_{t+1},\end{array}$$
而非这种方式：
$$\begin{array}{rl}{v}_{t+1}& =\mu \ast v_t+\text{lr}\ast g_{t+1},\\ p_{t+1}& =p_t-v_{t+1}.\end{array}$$
基于Pytorch源码，对自适应学习率的学习可见：基于Pytorch源码对自适应学习率进行学习

参考文档
各优化器详解，含公式、附图、优劣：深度学习——优化器算法Optimizer详解（BGD、SGD、MBGD、Momentum、NAG、Adagrad、Adadelta、RMSprop、Adam）
Pytorch的SGD类：torch.optim.sgd.py
Pytorch的Optimizer类：torch.optim.optimizer.py