zhengruifeng commented on a change in pull request #26739:
[SPARK-29967][ML][PYTHON] KMeans support instance weighting
URL: https://github.com/apache/spark/pull/26739#discussion_r353519210
##########
File path: mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
##########
@@ -278,30 +287,32 @@ class KMeans private (
val bcCenters = sc.broadcast(centers)
// Find the new centers
- val collected = data.mapPartitions { points =>
+ val collected = data.mapPartitions { pointsAndWeights =>
val thisCenters = bcCenters.value
val dims = thisCenters.head.vector.size
val sums = Array.fill(thisCenters.length)(Vectors.zeros(dims))
- val counts = Array.fill(thisCenters.length)(0L)
- points.foreach { point =>
- val (bestCenter, cost) =
distanceMeasureInstance.findClosest(thisCenters, point)
+ // clusterWeightSum is needed to calculate cluster center
+ // cluster center =
+ // sample1 * weight1/clusterWeightSum + sample2 *
weight2/clusterWeightSum + ...
+ val clusterWeightSum = Array.fill(thisCenters.length)(0.0)
+
+ pointsAndWeights.foreach { case (point, weight) =>
+ var (bestCenter, cost) =
distanceMeasureInstance.findClosest(thisCenters, point)
+ cost *= weight
costAccum.add(cost)
- distanceMeasureInstance.updateClusterSum(point, sums(bestCenter))
- counts(bestCenter) += 1
+ distanceMeasureInstance.updateClusterSum(point, sums(bestCenter),
weight)
+ clusterWeightSum(bestCenter) += weight
}
- counts.indices.filter(counts(_) > 0).map(j => (j, (sums(j),
counts(j)))).iterator
- }.reduceByKey { case ((sum1, count1), (sum2, count2)) =>
+ clusterWeightSum.indices.filter(clusterWeightSum(_) > 0)
+ .map(j => (j, (sums(j), clusterWeightSum(j)))).iterator
+ }.reduceByKey { case ((sum1, clusterWeightSum1), (sum2,
clusterWeightSum2)) =>
axpy(1.0, sum2, sum1)
- (sum1, count1 + count2)
+ (sum1, clusterWeightSum1 + clusterWeightSum2)
}.collectAsMap()
- if (iteration == 0) {
- instr.foreach(_.logNumExamples(collected.values.map(_._2).sum))
- }
-
Review comment:
1, I guess we need to add a new var `count: Long` to get the total count of
dataset, since in other algs like `LinearSVC`,`LogisticRegression`,
`instr.logNumExamples` logs the **unweighted** count;
2, Since more and more algs support weightCol, I think we may added a new
method like `instr.logSumOfWeights`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
