cuML: 快速GPU加速的机器学习库

cuML简介

cuML是RAPIDS项目的一部分,是一套基于GPU加速的机器学习算法库。它提供了与流行的scikit-learn库兼容的API,让数据科学家可以轻松地将现有的机器学习工作流程迁移到GPU上,而无需深入了解CUDA编程的细节。

cuML的主要特点包括:

• 与scikit-learn兼容的API,便于从CPU代码迁移
• 利用GPU加速算法,可以比CPU实现快10-50倍
• 支持单GPU、多GPU以及多节点多GPU的分布式计算
• 与其他RAPIDS项目如cuDF无缝集成

支持的算法

cuML支持广泛的机器学习算法,主要包括:

• 聚类:DBSCAN、K-Means等
• 降维:PCA、UMAP、t-SNE等
• 回归与分类:线性回归、逻辑回归、随机森林等
• 最近邻:KNN搜索与分类
• 时间序列:ARIMA、Holt-Winters指数平滑等
• 模型解释:SHAP值解释器

大多数算法都支持单GPU执行,部分算法如DBSCAN、K-Means等还支持多GPU和多节点分布式计算。

使用示例

以下是一个使用cuML进行DBSCAN聚类的简单示例:

import cudf
from cuml.cluster import DBSCAN

# 创建GPU DataFrame
gdf_float = cudf.DataFrame()
gdf_float['0'] = [1.0, 2.0, 5.0]
gdf_float['1'] = [4.0, 2.0, 1.0]
gdf_float['2'] = [4.0, 2.0, 1.0]

# 设置并拟合DBSCAN聚类
dbscan_float = DBSCAN(eps=1.0, min_samples=1)
dbscan_float.fit(gdf_float)

print(dbscan_float.labels_)

cuML还支持使用Dask进行多GPU和多节点计算。下面是一个使用Dask进行分布式KNN搜索的示例:

from dask_cuda import LocalCUDACluster
from dask.distributed import Client
import dask_cudf

# 创建本地CUDA集群
cluster = LocalCUDACluster(protocol="ucx", 
                           enable_tcp_over_ucx=True,
                           enable_nvlink=True)
client = Client(cluster)

# 并行读取CSV文件
df = dask_cudf.read_csv("/path/to/csv")

# 拟合KNN模型并查询
from cuml.dask.neighbors import NearestNeighbors
nn = NearestNeighbors(n_neighbors=10, client=client)
nn.fit(df)
neighbors = nn.kneighbors(df)

安装与使用

cuML可以通过conda或pip安装:

conda install -c rapidsai -c conda-forge cuml cudatoolkit=11.5

或

pip install cuml-cu11 --extra-index-url=https://pypi.ngc.nvidia.com

cuML需要CUDA环境支持,推荐使用Linux系统。从23.10版本开始,cuML还提供了CPU执行功能,可以在不支持GPU的系统上使用部分功能。

性能优势

对于大型数据集,cuML的GPU实现可以比CPU版本快10-50倍。例如在MNIST数据集上,cuML的PCA算法比scikit-learn快约30倍。具体的性能对比可以参考cuML提供的基准测试笔记本。

总结

cuML为数据科学家提供了一种简单的方式来利用GPU加速机器学习工作流程。通过与scikit-learn兼容的API,用户可以轻松迁移现有代码,同时获得显著的性能提升。随着对大规模数据集分析需求的增长,cuML将成为数据科学家工具箱中越来越重要的一员。