文章来源于网络收集而来，版权归原创者所有，如有侵权请及时联系！

Caffe 代码阅读 - Solver

发布于 2025-02-25 23:04:57 字数 4440 浏览 0 评论 0 收藏 0

前面我们聊了 Net 组装的内容，接下来我们来看看 Solver 的内容。Solver 主体有两部分：初始化和训练。初始化内容相对比较简单，这里就不说了；下面我们来说说训练中的几个关键函数。

核心函数：Step

真正的训练在 Step 函数内，这里有多卡训练的关键回调函数：on_start() 和 on_gradient_ready()，具体的调用方法我们后面再说，在这两个回调函数中间有两个重要的过程：ForwardBackward 和 UpdateSmoothedLoss。在 on_gradient_ready 之后有一个关键函数 ApplyUpdate()，这里面的代码在 Sgd_solver 中。下面我们详细看一下。

ForwardBackward

这里主要调用了 Net 中的代码，主要完成了前向后向的计算，前向用于计算模型的最终输出和 Loss，后向用于计算每一层网络和参数的梯度。对于前向后向的具体内容这里需要详细叙述了，唯一值得一提的是前向的 Loss 计算，这部分代码实际上实在 Layer 里面，具体涉及到 loss_weight 这个参数相关的初始化和 loss() 的判断，同时还有 Loss_Layer 在 Setup 函数中的初始化。

UpdateSmoothedLoss

这个函数主要做 Loss 的平滑。由于 Caffe 的训练方式是 SGD，我们无法把所有的数据同时放入模型进行训练，那么部分数据产生的 Loss 就可能会和全样本的平均 Loss 不同，在必要时候将 Loss 和历史过程中更新的 Loss 求平均就可以减少 Loss 的震荡问题。代码中的平滑方法比较简单，大家一看便知。

下面就是 ApplyUpdate 函数，这个函数真正完成了参数更新的任务。Caffe 的参数更新只利用了模型的梯度信息，没有利用二阶信息。下面就详细介绍下更新参数的几个过程：

GetLearningRate
ClipGradients
Normalize
Regularize
ComputeUpdateValue

GetLearningRate

learning rate 的故事我们前面已经聊过了，在 CNN 训练中这确实是个大问题。Caffe 为了让 learning rate 的设计更灵活，提供了一系列的 learning rate 方案：

fixed ：lr 永远不变
step ： $lr=baselr*gamma^{iter / stepsize}$
exp ： $lr=baselr*gamma^{iter}$
inv ： $lr=baselr*(1+gamma*iter)^{-power}$
multistep ：直接写 iter 在某个范围内时 lr 应该是多少
poly ： $lr=baselr*(1-\frac{iter}{maxiter})^{power}$
sigmoid ： $lr=baselr*\frac{1}{1+e^{-gamma*(iter-stepsize)}}$

这些方案各有优劣，选择自己顺手的就好。

ClipGradients

这一步主要是对梯度值做一个限制，如果梯度值过大，那么这里就会对梯度做一个修剪，对所有的参数乘以一个缩放因子，使得所有参数的平方和不超过参数中设定的梯度总值。这个功能感觉上像是对全局函数设置了一个 Trust Region，可以防止更新的量过大二导致梯度发散。我认为这一步的想法是很好的，但是实际操作中可能会有问题。实际中可能只有部分参数的梯度比较大，而其他参数的梯度本身比较小，那么对所有的参数乘以相同的因子会让一些本来比较小的参数变得更小，这样会带来一些不公平。

Normalize

这一步同样考虑了一些单一 Batch 不足以完成训练的问题，通过限制每个 Batch 的更新量来控制更新总量，代码比较简单。

Regularize

到这一步终于要计算正则项的梯度了。Caffe 提供两种正则方法 - L2 和 L1，其中 L2 采用了标准的梯度下降方法，L1 采用了 sub-gradient 的计算方法。L2 的优化计算比较简单，没有什么好说的，但是 L1 的计算还是有点值得玩味的地方的。这里采用的 sub-gradient 方法其实本身没有什么问题，不过 Lasso 的优化还可以有其他的方法，这个问题以后可以再细聊。

ComputeUpdateValue

到这里，我们终于来到了梯度计算的最后一站，这时候我们终于完成了对梯度的计算，下面该考虑 lr 和梯度结合起来如何计算最终的梯度优化值了。sgd 方法主要采用 momentum 加梯度的优化方法。关于 momentum 的优势我们前面已经聊过了。除此之外，Caffe 还提供了一系列的梯度计算方法，这些优化方法各有特点，以后我们可以慢慢来看。

当计算完这一步，我们就可以调用 Blob 中的 Update 把每个参数的 data 和 diff 进行相加，计算出最终的结果。这样，整个优化过程就完成了。至于剩下的一些内容都不是核心过程，就略去不看了。

如果我们采用单卡训练的策略，那么阅读代码到这里也差不多了。不过多卡训练对于大规模的训练任务来说是必不可少的，所以我们接下来趁热打铁地看看 Caffe 的多卡训练。

多卡训练算法

Caffe 的多卡训练算法总体思路是数据并行，我们用不同的 GPU 处理不同的数据，然后将所有的梯度更新汇总。由于 Solver 在训练中给了两个回调函数，多卡训练也主要利用了这两个回调函数进行：

on_start()：将参数拷贝到每一个 GPU 中。
ForwardBackward()：每个 GPU 各自计算自己的前向后向结果。
on_gradient_ready()：将反向梯度汇总到一起。
ApplyUpdate()：在汇总的线程上进行参数更新

其中第 2 步由每一个 CPU 线程和自己的 GPU 并行完成，第 4 步由汇总的 CPU 和自己的 GPU 完成，剩下的 1，3 两步主要是完成数据传输的任务，也是多卡计算中主要完成的部分。

Caffe 采用树型结构进行参数传递，其中一个 CPU 线程和 GPU 作为树型结构的根，其他的则作为根下面的节点。为了更快地传输 GPU 数据，树型结构的构建要考虑 GPU 之间是否相近，比方说两个 GPU 之间是否可以进行 P2P 的直传。在前面的翻译博客中我们已经聊过 GPU 之间数据传输的问题了，这里的树形结构也主要以此做考虑。

我们假设 4 块 GPU 的拓扑结构如下：

nvidia-smi topo -m
       GPU0   GPU1   GPU2   GPU3  
GPU0   X     PHB    SOC    SOC    
GPU1   PHB    X     SOC    SOC    
GPU2   SOC    SOC    X     PHB   
GPU3   SOC    SOC    PHB    X

那么我们构造出的树型结构如下所示，数据传输也是按照这样的结构传输：