一文搞懂网络训练中一个epoch中到底做了什么+zero_grad可以放在backward后面吗(train.py)

建议阅读时间：7分钟

参考：

• d2lhttps://zh-v2.d2l.ai/
• 博客https://blog.csdn.net/PanYHHH/article/details/107361827

---

前言了解：

1. model/net会存储两个矩阵，分别是params矩阵用于存储权重参数；以及params.grad用于存储梯度参数。其中最吃GPU计算资源的是params.grad矩阵的计算。
2. optimizer是一个可以放在CPU上工作的不吃什么计算资源的object，它只是基于梯度下降算法，来更新model/net中params矩阵的“优化器”。硬件角度来理解：它只需要一个乘法运算器+一个加法运算器便可以组成。就是干下面这个的：

一个epoch到底做了什么

'''
loss 	:= cost function
net 	:= model
trainer	:= optimizer
'''
num_epochs = 3
for epoch in range(num_epochs):
	for X, y in data_iter:
		l = loss(net(X) ,y)
		trainer.zero_grad()
		l.backward()
		trainer.step()
	l = loss(net(features), labels)
	print(f'epoch {epoch + 1}, loss {l:f}')

为了很好的理解上面代码首先我们需要知道，在网络进行训练的过程中，我们会存储两个矩阵：分别是params矩阵用于存储权重参数；以及params.grad用于存储梯度参数。下面我们来将上面的网络过程进行数理：

1. 取数据

for X, y in data_iter这句话用来取数据，Dataloader作为一个迭代器，他每次调用batch_size次Dataset中的__getitem__方法或者IterableDataset中的__iter__方法，来取batch_size对数据(包括signal/input和target/label)。

2. 正向传播

l = loss(net(X), y)这一句是将样本输入网络，与权重相乘后输出$\hat{y}$，将$\hat{y}$与y传入cost function进行loss计算。即整行是完成了一次完整的正向传播过程，在这个过程中会生成计算图用于后续的梯度计算。

3. 梯度清零

trainer.zero_grad()用于将记录梯度的矩阵清零。为了更好的理解这句话，我们来看一下它的代码实现：

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

4. 梯度更新

l.backward()用于进行反向传播计算梯度并将梯度值存储在param.grad中，即更新了梯度矩阵。

5. 更新参数

trainer.step()用于更新参数值，即更新了params这个权重矩阵，在最常用的mini batch SGD中，他的算法如下：$params=params-\frac{lr\ast grad}{batchsize}$。

下面显示了torch.optim.SGD().step()源码：

    def step(self, closure=None):
        """Performs a single optimization step.
        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()
 
        for group in self.param_groups:
            weight_decay = group['weight_decay']
            momentum = group['momentum']
            dampening = group['dampening']
            nesterov = group['nesterov']
 
            for p in group['params']:
                if p.grad is None:
                    continue
                d_p = p.grad.data
                if weight_decay != 0:
                    d_p.add_(weight_decay, p.data)
                if momentum != 0:
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
                    else:
                        buf = param_state['momentum_buffer']
                        buf.mul_(momentum).add_(1 - dampening, d_p)
                    if nesterov:
                        d_p = d_p.add(momentum, buf)
                    else:
                        d_p = buf
 
                p.data.add_(-group['lr'], d_p)
 
        return loss

问题与解答

1. 可以最后梯度清零吗？

当然你也可以选择在最后梯度清零，但可能会污染第一轮的梯度。