1. 首页
  2. 开发者
  3. 机器学习

Caffe代码阅读——Solver

Caffe代码阅读——Solver

      前面我们聊了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永远不变
  • stepCaffe代码阅读——Solver
  • expCaffe代码阅读——Solver
  • invCaffe代码阅读——Solver
  • multistep:直接写iter在某个范围内时lr应该是多少
  • polyCaffe代码阅读——Solver
  • sigmoidCaffe代码阅读——Solver

 

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

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在训练中给了两个回调函数,多卡训练也主要利用了这两个回调函数进行:

  1. on_start():将参数拷贝到每一个GPU中。
  2. ForwardBackward():每个GPU各自计算自己的前向后向结果。
  3. on_gradient_ready():将反向梯度汇总到一起。
  4. 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

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

Caffe代码阅读——Solver

      这样1,3的数据传递就解决了,具体的过程请详细阅读代码,这里就不叙述了。

      对Caffe代码的基本介绍就到这里了,我们对代码的整体结构有了比较清晰的认识,下面我们将分析模型中各个部分的特性。

发布者: aihot,转转请注明出处:http://www.aiwuyun.net/archives/3901.html

发表评论

登录后才能评论

联系我们

在线咨询:点击这里给我发消息

邮件:admin@aiwuyun.net

工作时间:周一至周五,9:30-18:30,节假日休息