分布式FTRL优化算法的实现

之前实现过Parameter Server框架下的分布式FTRL优化算法,用的是DMLC的ps-lite。在PS架构下,集群分为worker、server、scheduler三种线程。其中worker负责梯度的计算,server负责参数的更新和分布式存储,scheduler负责集群节点的管理和状态监控。将代码实现简单介绍如下:

        1.  在main.cpp中完成启三个线程的启动过程。

         2. 在worker.cpp中完成梯度的计算并向server 推送算得的梯度。其中,(1)oneline_learning_threadpool是做在线学习用的,数据流主要是从kafka中读取,所有数据只训练一遍;batch_learning_threadpool是做批量学习用的,数据主要是从磁盘或者HDFS中读取,训练数据可以反复训练多次。(2)calculate_batch_gradient_threadpool负责根据样本计算梯度,在计算过程中worker和server通过ps-lite提供的pull/push接口拉取模型参数和推送样本梯度值。针对梯度计算,过程如下:首先,读取一批训练数据,代码328行至339行得到该批次数据的所有特征ID:all_keys,340行至343行进行过滤得到唯一的特征ID,348行从server中拉取所需权重,且权重是按照特征ID值大小排序的。352行至363行,根据pull得到的特征权重,计算该批次中每条样本的w*x乘积,这里有个巧妙的方法可以在O(m)时间复杂度完成计算,最开始的做法是把pull回来的数据放到map中,这样需要m*O(logn)时间复杂度。经过试验对比,O(m)时间复杂度的算法速度快了很多。我曾经尝试过其它的hashmap来替代std::unordered_map,效果都没有明显提升。

        3. 在server.cpp中完成模型的更新,根据FTRL公式,server收到worker推送过来的梯度之后进行一系列的计算更新z和w,从而完成模型的更新。156行的Push函数是ps-lite会调用的回调函数,也是用户实现模型更新逻辑的地方。

        4. dump.cpp函数将二进制的模型文件转成明文。

        5. 在io文件夹中,io.h是所有io接口的声明。load_data_from_local.cc从磁盘读取数据,load_data_from_kafka.cc从kafka中读取流式数据。其中load_minibatch_hash_data_fread函数采用fread函数读取数据,比用std::cin读取一行之后再分割的方法,在速度上快了近8倍。

        代码地址在:https://github.com/xswang/Field-aware-Factorization-Machine-ps ,之前是在三台机器上测试过分布式版本的,目前修改成单机版,如果要改成多机版也非常方便。

Leave a Reply

Your email address will not be published. Required fields are marked *