Spark RAPID ML 是一个开源 Python 包,它可以使 NVIDIA GPU 加速 PySpark MLlib。它提供了与 PySpark MLlib DataFrame API 兼容,并在使用支持的算法进行训练时加速。想要了解更多信息,请查看 新的 GPU 库降低 Apache Spark ML 的计算成本。
PySpark MLlib DataFrame API 的兼容性意味着它可以更容易地融入现有的 PySpark ML 应用程序,最多只需更改包导入。K-means 算法如下所示。更改包导入是使用此库启用 GPU 加速所需的唯一额外步骤。
PySpark MLlib
from pyspark.ml.clustering import KMeans kmeans_estm = KMeans()\ .setK(100)\ .setFeaturesCol("features")\ .setMaxIter(30) kmeans_model = kmeans_estm.fit(pyspark_data_frame) kmeans_model.write().save("saved-model") transformed = kmeans_model.transform(pyspark_data_frame) |
Spark RAPID ML
from spark_rapids_ml.clustering import KMeans kmeans_estm = KMeans()\ .setK(100)\ .setFeaturesCol("features")\ .setMaxIter(30) kmeans_model = kmeans_estm.fit(pyspark_data_frame) kmeans_model.write().save("saved-model") transformed = kmeans_model.transform(pyspark_data_frame) |
在 GPU 加速的 Databricks 的 AWS 托管