分布式机器学习:逻辑回归的并行化实现(PySpark)
算法的完整实现代码我已经上传到了GitHub仓库:Distributed-ML-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\}\)。则逻辑回归的损失函数为:
这里
写成这个形式就已经可以用诸如Pytorch这类工具来进行自动求导然后采用梯度下降法求解了。不过若需要用表达式直接计算出梯度,我们还需要将损失函数继续化简为:
可将梯度表示如下:
2 基于Spark的并行化实现
逻辑回归的目标函数常采用梯度下降法求解,该算法的并行化可以采用如下的Map-Reduce架构(图片来自王树森老师的YouTube课程并行计算与机器学习(1/3)[2]):
先将第\(t\)轮迭代的权重广播到各worker,各worker计算一个局部梯度(map过程),然后再将每个节点的梯度聚合(reduce过程),最终对参数进行更新。
在Spark中每个task对应一个分区,决定了计算的并行度(分区的概念详间我们上一篇博客Spark: 单词计数(Word Count)的MapReduce实现(Java/Python) )。我们假设共有3个分区,则在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的完整实现代码如下:
'''
Descripttion:
Version: 1.0
Author: ZhangHongYu
Date: 2022-05-26 21:02:38
LastEditors: ZhangHongYu
LastEditTime: 2022-07-01 16:22:53
'''
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
import matplotlib.pyplot as plt
import sys
import os
os.environ['PYSPARK_PYTHON'] = sys.executable
n_threads = 3 # Number of local threads
n_iterations = 1500 # Number of iterations
eta = 0.1 # 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
def draw_acc_plot(accs, n_iterations):
def ewma_smooth(accs, alpha=0.9):
s_accs = np.zeros(n_iterations)
for idx, acc in enumerate(accs):
if idx == 0:
s_accs[idx] = acc
else:
s_accs[idx] = alpha * s_accs[idx-1] + (1 - alpha) * acc
return s_accs
s_accs = ewma_smooth(accs, alpha=0.9)
plt.plot(np.arange(1, n_iterations + 1), accs, color="C0", alpha=0.3)
plt.plot(np.arange(1, n_iterations + 1), s_accs, color="C0")
plt.title(label="Accuracy on test dataset")
plt.xlabel("Round")
plt.ylabel("Accuracy")
plt.savefig("logistic_regression_acc_plot.png")
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")\
.master("local[%d]" % n_threads)\
.getOrCreate()
matrix = np.concatenate(
[X_train, np.ones((n_train, 1)), y_train.reshape(-1, 1)], axis=1)
points = spark.sparkContext.parallelize(matrix).cache()
# Initialize w to a random value
w = 2 * np.random.ranf(size=D + 1) - 1
print("Initial w: " + str(w))
accs = []
for t in range(n_iterations):
print("On iteration %d" % (t + 1))
g = points.map(lambda point: gradient(point, w)).reduce(add)
w -= eta * 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)
accs.append(acc)
print("iterations: %d, accuracy: %f" % (t, acc))
print("Final w: %s " % w)
print("Final acc: %f" % acc)
spark.stop()
draw_acc_plot(accs, n_iterations)
注意.master("local[%d]" % n_threads)
中的n_threads
是我们在本地单机多线程调试模式下所设置的线程数,也就是Spark中的默认分区数,我们此处将n_threads
设置为3,则Spark就会默认划分出3个分区。我们在代码中采用breast cancer数据集进行训练和测试,该数据集是个二分类数据集。模型初始权重采用随机初始化。
最后,我们来看一下算法的输出结果。
初始权重如下:
Initial w: [ 0.20733249 -0.68270323 -0.23539134 0.46125717 -0.27736064 -0.36072597
0.92549048 -0.18432978 0.77991313 0.54054734 0.48559498 -0.23031733
0.67125099 0.57301018 0.69243332 -0.4791771 -0.76039149 0.15924619
0.01321836 -0.19976038 0.576716 0.50379885 0.58670905 -0.38590575
-0.48719581 -0.91967718 0.73359703 0.28669715 0.56688998 0.97444464
-0.44361797]
最终的模型权重与在测试集上的准确率结果如下:
Final w: [ 1.15974825e+04 1.29973800e+04 6.52553107e+04 2.10241061e+04
8.86514067e+01 -1.10587723e+02 -2.97300359e+02 -1.27131718e+02
1.59369309e+02 7.84967515e+01 -4.03071071e+01 8.13799814e+02
-1.30662140e+03 -4.04474691e+04 5.34490109e+00 -2.28709226e+01
-4.24236287e+01 -8.04493849e+00 1.12580376e+01 7.93202730e-01
1.25640151e+04 1.51951403e+04 6.46383775e+04 -3.18968898e+04
9.95884228e+01 -4.01750499e+02 -6.93005010e+02 -1.78725566e+02
1.93133380e+02 6.01062122e+01 1.52932953e+03]
Final acc: 0.947368
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()
操作来替代reduce()
操作,该操作使用树形聚合方法来减少网络和计算延迟。
treeAggregate()
函数原型如下:
RDD.treeAggregate(zeroValue, seqOp, combOp, depth=2)
其中zeroValue
为聚合结果的初始值,seqOp
函数用于定义单分区(partition)做聚合操作的方法,它第一个参数为聚合结果,第二个参数为分区中的数据变量。combOp
定义对分区之间做聚合的方法,它第一个参数为第二个参数都为聚合结果。
depth
为聚合树的深度。
此处我们的聚合操作比较简单,聚合结果初始值设置为0.0
,seqOp
和combOp
都设置为add
算子即可:
g = points.map(lambda point: gradient(point, w_br.value))\
.treeAggregate(0.0, add, add)
6 算法收敛性和复杂度分析
6.1 收敛性和计算复杂度
算法的ACC曲线如下图所示:
可见我们的算法总体呈现收敛。
这里的损失函数\(l( \space \cdot \space )\)是光滑凸函数(非强凸函数,它只在局部呈现强凸性),如我们在博客《数值优化:经典一阶确定性算法及其收敛性分析》中所分析,假设目标函数\(f: \mathbb{R}^d \rightarrow \mathbb{R}\)是凸函数,且\(\beta\)-光滑,当步长\(\eta = \frac{1}{\beta}\)时,梯度下降法具有\(\mathcal{O}(\frac{1}{t})\)的次线性收敛速率:
这意味着在梯度下降求解逻辑回归问题的迭代次数复杂度为\(\mathcal{O}(\frac{1}{\varepsilon})\),也即\(\mathcal{O}(\frac{1}{\varepsilon})\)轮后会取得\(\varepsilon\)的精度。
尽管梯度的计算可以被分摊到个计算节点上,然而梯度下降的迭代是串行的。每轮迭代中,Spark会执行同步屏障(synchronization barrier)来确保在各worker开始下一轮迭代前\(w\)已被更新完毕。如果存在掉队者(stragglers),其它worker就会空闲(idle)等待,直到下一轮迭代。故相比梯度的计算,其迭代计算的“深度”(depth)是其计算瓶颈。
6.2 通信复杂度
map过程显然是并行的,并不需要通信。broadcast过程需要一对多通信,并且reduce过程需要多对一通信(都按照树形结构)。故对于每轮迭代,总通信时间按
增长。
这里\(p\)为除去driver节点的运算节点个数,\(L\)是节点之间的通信延迟。\(B\)是节点之间的通信带宽。\(M\)是每轮通信中节点间传输的信息大小。故消息能够够以\(\mathcal{O}(\log p)\)的通信轮数在所有节点间传递。
参考
- [1] GiHub: Spark官方Python样例
- [2] 王树森YouTube课程:并行计算与机器学习(1/3)
- [3] 刘铁岩,陈薇等. 分布式机器学习:算法、理论与时间[M]. 机械工业出版社, 2018.
- [4] 许利杰,方亚芬. 大数据处理框架Apache Spark设计与实现[M]. 电子工业出版社, 2021.
- [5] Stanford CME 323: Distributed Algorithms and Optimization (Lecture 7)