分布式机器学习:逻辑回归的并行化实现(PySpark)

1. 梯度计算式导出

我们在博客《统计学习:逻辑回归与交叉熵损失(Pytorch实现)》中提到,设\(w\)为权值(最后一维为偏置),样本总数为\(N\)\(\{(x_i, y_i)\}_{i=1}^N\)为训练样本集。样本维度为\(D\)\(x_i\in \mathbb{R}^{D+1}\)(最后一维扩充),\(y_i\in\{0, 1\}\)。则逻辑回归的损失函数为:

\[\mathcal{l}(w) = \sum_{i=1}^{N}\left[y_{i} \log \pi_{w}\left(x_{i}\right)+\left(1-y_{i}\right) \log \left(1-\pi_w\left(x_{i}\right)\right)\right] \]

这里

\[\begin{aligned} \pi_w(x) = p(y=1 \mid x; w) =\frac{1}{1+\exp \left(-w^{T} x\right)} \end{aligned} \]

写成这个形式就已经可以用诸如Pytorch这类工具来进行自动求导然后采用梯度下降法求解了。不过若需要用表达式直接计算出梯度,我们还需要将损失函数继续化简为:

\[\mathcal{l}(w) = -\sum_{i=1}^N(y_i w^T x_i - \log(1 + \exp(w^T x_i))) \]

可将梯度表示如下:

\[\nabla_w{\mathcal{l}(w)} = -\sum_{i=1}^N(y_i - \frac{1}{\exp(-w^Tx)+1})x_i \]

2. 基于Spark的并行化实现

逻辑回归的目标函数常采用梯度下降法求解,该算法的并行化可以采用如下的Map-Reduce架构:

先将第\(t\)轮迭代的权重广播到各worker,各worker计算一个局部梯度(map过程),然后再将每个节点的梯度聚合(reduce过程),最终对参数进行更新。

在Spark中每个task对应一个分区,决定了计算的并行度(分区的概念详间我们上一篇博客Spark: 单词计数(Word Count)的MapReduce实现(Java/Python) )。在Spark的实现过程如下:

  • map阶段: 各task运行map()函数对每个样本\((x_i, y_i)\)计算梯度\(g_i\), 然后对每个样本对应的梯度运行进行本地聚合,以减少后面的数据传输量。如第1个task执行reduce()操作得到\(\widetilde{g}_1 = \sum_{i=1}^3 g_i\) 如下图所示:

  • reduce阶段:使用reduce()将所有task的计算结果收集到Driver端进行聚合,然后进行参数更新。

在上图中,训练数据用points:PrallelCollectionRDD来表示,参数向量用\(w\)来表示,注意参数向量不是RDD,只是一个单独的参与运算的变量。

此外需要注意一点,虽然每个task在本地进行了局部聚合,但如果task过多且每个task本地聚合后的结果(单个gradient)过大那么统一传递到Driver端仍然会造成单点的网络平均等问题。为了解决这个问题,Spark设计了性能更好的treeAggregate()操作,使用树形聚合方法来减少网络和计算延迟,我们在第5部分会详细叙述。

3. PySpark实现代码

PySpark的完整实现代码如下:

from sklearn.datasets import load_breast_cancer
import numpy as np
from pyspark.sql import SparkSession
from operator import add
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

n_slices = 3  # Number of Slices
n_iterations = 300  # Number of iterations
alpha = 0.01  # iteration step_size


def logistic_f(x, w):
    return 1 / (np.exp(-x.dot(w)) + 1)


def gradient(point: np.ndarray, w: np.ndarray) -> np.ndarray:
    """ Compute linear regression gradient for a matrix of data points
    """
    y = point[-1]    # point label
    x = point[:-1]   # point coordinate
    # For each point (x, y), compute gradient function, then sum these up
    return - (y - logistic_f(x, w)) * x


if __name__ == "__main__":

    X, y = load_breast_cancer(return_X_y=True)

    D = X.shape[1]
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=0)
    n_train, n_test = X_train.shape[0], X_test.shape[0]

    spark = SparkSession\
        .builder\
        .appName("Logistic Regression")\
        .getOrCreate()

    matrix = np.concatenate(
        [X_train, np.ones((n_train, 1)), y_train.reshape(-1, 1)], axis=1)

    points = spark.sparkContext.parallelize(matrix, n_slices).cache()

    # Initialize w to a random value
    w = 2 * np.random.ranf(size=D + 1) - 1
    print("Initial w: " + str(w))

    for t in range(n_iterations):
        print("On iteration %d" % (t + 1))
        g = points.map(lambda point: gradient(point, w)).reduce(add)
        w -= alpha * g

        y_pred = logistic_f(np.concatenate(
            [X_test, np.ones((n_test, 1))], axis=1), w)
        pred_label = np.where(y_pred < 0.5, 0, 1)
        acc = accuracy_score(y_test, pred_label)
        print("iterations: %d, accuracy: %f" % (t, acc))

    print("Final w: %s " % w)
    print("Final acc: %f" % acc)

    spark.stop()

注意spark.sparkContext.parallelize(matrix, n_slices)中的n_slices就是Spark中的分区数。我们在代码中采用breast cancer数据集进行训练和测试,该数据集是个二分类数据集。模型初始权重采用随机初始化。

最后,我们来看一下算法的输出结果。

初始权重如下:

Initial w: [-0.0575882   0.79680833  0.96928013  0.98983501 -0.59487909 -0.23279241
 -0.34157571  0.93084048 -0.10126002  0.19124314  0.7163746  -0.49597826
 -0.50197367  0.81784642  0.96319482  0.06248513 -0.46138666  0.76500396
  0.30422518 -0.21588114 -0.90260279 -0.07102884 -0.98577817 -0.09454256
  0.07157487  0.9879555   0.36608845 -0.9740067   0.69620032 -0.97704433
 -0.30932467]

最终的模型权重与在测试集上的准确率结果如下:

Final w: [ 8.22414803e+02  1.48384087e+03  4.97062125e+03  4.47845441e+03
  7.71390166e+00  1.21510016e+00 -7.67338147e+00 -2.54147183e+00
  1.55496346e+01  6.52930570e+00  2.02480712e+00  1.09860082e+02
 -8.82480263e+00 -2.32991671e+03  1.61742379e+00  8.57741145e-01
  1.30270454e-01  1.16399854e+00  2.09101988e+00  5.30845885e-02
  8.28547658e+02  1.90597805e+03  4.93391021e+03 -4.69112527e+03
  1.10030574e+01  1.49957834e+00 -1.02290791e+01 -3.11020744e+00
  2.37012097e+01  5.97116694e+00  1.03680530e+02] 
Final acc: 0.923977

可见我们的算法收敛良好。

4.关于冗余存储的反思

注意根据我们以上的代码实现中的

map(lambda point: gradient(point, w)).reduce(add)

这一行中,我们求梯度的函数gradient会根据w将每一个训练样本点map到其对应梯度值。w的拷贝会被发送给每个计算节点的每个CPU。比如,假设我们有一个4个CPU的计算节点。

默认当map过程发生时,所有被map过程需要的数据会被发往mapper,而此处每个CPU都有一个mapper,故如果该计算节点有4个CPU,我们实际上会发送4个w的拷贝到该节点,如下图所示:

之所以会这样,是因为此处假定w会被修改,必须为每个CPU单独存储w拷贝以解决并发写的问题。然而,当我们计算每一步的梯度时,w并未被修改,故此处不存在并发写的问题。这导致我们浪费了存储空间,因为本可将w存储在各个节点的共享内存中的。

为了解决此问题,我们可以将w进行广播,这样它只会被发到每个计算节点一次(而不是每个CPU一次)。为了实现这个想法,我们将w定义为一个广播变量来使用,如下面代码所示:

# Initialize w to a random value
w = 2 * np.random.ranf(size=D + 1) - 1
print("Initial w: " + str(w))

for t in range(n_iterations):
    print("On iteration %d" % (t + 1))
    w_br = spark.sparkContext.broadcast(w)
    
    g = points.map(lambda point: gradient(point, w_br.value)).reduce(add)

    w -= alpha * g

当我们初始化w时,我们首先将其声明为一个广播变量。在每一轮梯度下降的迭代中,我们需要引用w的值。最后,我们在w被更新后重新广播w。这样,w在每个机器上被高效地存储(每个机器一份,而不是多份)。

5.关于聚合效率的反思

正如我们前面所说,我们可以用性能更好的treeAggregate()操作,使用树形聚合方法来减少网络和计算延迟。
treeAggregate()函数原型如下:

RDD.treeAggregate(zeroValue, seqOp, combOp, depth=2)

其中zeroValue为聚合结果的初始值,seqOp函数用于定义单分区(partition)做聚合操作的方法,它第一个参数为聚合结果,第二个参数为分区中的数据变量。combOp定义对分区之间做聚合的方法,它第一个参数为第二个参数都为聚合结果。
depth为聚合树的深度。

此处我们的聚合操作比较简单,聚合结果初始值设置为0.0seqOpcombOp都设置为add算子即可:

g = points.map(lambda point: gradient(point, w_br.value))\
    .treeAggregate(0.0, add, add)

6. 复杂度和通信时间分析

6.1 复杂度

如我们在博客《数值优化:经典一阶确定性算法及其收敛性分析》中所分析,梯度下降法在光滑强凸条件下有拥有线性收敛速率,其迭代次数复杂度为\(\mathcal{O}(\text{log}\frac{1}{\varepsilon})\)

尽管梯度的计算可以被分摊到个计算节点上,然而梯度下降的迭代是串行的。每轮迭代中,Spark会执行同步屏障(synchronization barrier)来确保在各worker开始下一轮迭代前w已被更新完毕。如果存在掉队者(stragglers),其它worker就会空闲(idle)等待,直到下一轮迭代。故相比梯度的计算,其迭代计算的“深度”(depth)是其计算瓶颈。

6.2 通信时间

map过程显然是并行的,并不需要通信。broadcast过程需要一对多通信,并且reduce过程需要多对一通信(都按照树形结构)。故对于每轮迭代,总通信时间按

\[2\text{log}_2(p)(L + \frac{m}{B}) \]

增长。
这里\(p\)为除去driver节点的运算节点个数,\(L\)是节点之间的通信延迟。\(B\)是节点之间的通信带宽。\(M\)是每轮通信中节点间传输的信息大小。

参考

posted @ 2022-05-27 19:00  orion-orion  阅读(160)  评论(0编辑  收藏  举报