pyspark中使用累加器Accumulator统计指标

评价分类模型的性能时需要用到以下四个指标

最开始使用以下代码计算，发现代码需要跑近一个小时，而且这一个小时都花在这四行代码上

# evaluate model

TP = labelAndPreds.filter(lambda (v, p): (v == 1 and p == 1)).count()

FP = labelAndPreds.filter(lambda (v, p): (v == 0 and p == 1)).count()

TN = labelAndPreds.filter(lambda (v, p): (v == 0 and p == 0)).count()

FN = labelAndPreds.filter(lambda (v, p): (v == 1 and p == 0)).count()

心想着理论上可以只扫描一遍数据就可以计算出这四个指标。

一开始在foreach函数中传递一个自定义评估函数，这个函数来统计上面四个指标，然后在函数里再使用全局变量TP，TN等。

但是程序跑完四个指标都还是0，跟初始化时候的一样。后来查资料，发现pyspark有Accumulator(累加器)可以解决这个问题。

代码如下：

# evaluate model

TP = sc.accumulator(0)　　#一开始直接用的TP = 0

FP = sc.accumulator(0)

TN = sc.accumulator(0)

FN = sc.accumulator(0)

def assess(v, p):

    global TP

    global FP

    global TN

    global FN

    #print 'tgl\t',v,p

    if(v == 1 and p == 1):

        TP += 1

    if(v == 0 and p == 1):

        FP += 1

    if(v == 0 and p == 0):

        TN += 1

    if(v == 1 and p == 0):

        FN += 1

print 'assess model %s' % time.ctime()

labelAndPreds.foreach(lambda(v,p): assess(v, p))

print "TP=", TP

print "FP=", FP

print "TN=", TN

print "FN=", FN

if (TP.value + FP.value) != 0:

      print "The precision = " + str(TP.value*1.0 / (TP.value+FP.value))

if (TP.value + FN.value) != 0:

      print "The recall = " + str(TP.value*1.0 / (TP.value+FN.value))

ps:

pyspark官方文档

[http://spark.apache.org/docs/latest/api/python/pyspark.html?highlight=accumulator#pyspark.Accumulator]