cuML - GPU 机器学习算法库

cuML 是一套实现机器学习算法和数学原语函数的库，其 API 与其它 RAPIDS 项目兼容。

cuML 使数据科学家、研究人员和软件工程师能够在 GPU 上运行传统的表格型机器学习任务，而无需深入了解 CUDA 编程。在大多数情况下，cuML 的 Python API 与 scikit-learn 的 API 相匹配。

对于大型数据集，这些基于 GPU 的实现比对应的 CPU 版本快 10-50 倍。性能详情可参考 cuML 基准测试 Notebook。

以下示例展示了如何通过 cuDF 在 GPU 上加载输入数据并计算 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_float = DBSCAN(eps=1.0, min_samples=1)
dbscan_float.fit(gdf_float)

print(dbscan_float.labels_)

输出：

0    0
1    1
2    2
dtype: int32

cuML 还支持多 GPU 和多节点多 GPU 操作（基于 Dask），且支持的算法列表持续增长。以下 Python 代码示例从 CSV 文件读取数据，并在单节点多 GPU 的 Dask 集群上执行最近邻（NearestNeighbors）查询：

初始化一个配置了 UCXX 的 LocalCUDACluster，用于快速传输 CUDA 数组：

from dask_cuda import LocalCUDACluster

# 创建单节点 CUDA 集群，每个设备一个 worker
cluster = LocalCUDACluster(protocol="ucx",
                           enable_tcp_over_ucx=True,
                           enable_nvlink=True,
                           enable_infiniband=False)

加载数据并执行 k-最近邻搜索。cuml.dask 的估计器也支持 Dask.Array 作为输入：

from dask.distributed import Client
client = Client(cluster)

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

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

更多示例请浏览我们的完整 API 文档 或查看示例 教程 Notebook。此外，你还可以在 notebooks-contrib 仓库 中找到完整的端到端示例。

支持的算法

安装

请参考 RAPIDS 版本选择器 获取通过 conda、pip 或 Docker 安装 nightly 版或正式版 cuML 包的命令行。

从源码构建/安装

请参考构建指南。

Scikit-learn 兼容性

cuML 兼容 scikit-learn 1.4 及以上版本。

模型序列化与安全

cuML 模型可以使用 pickle 或 joblib 序列化，并在之后加载用于推理。cuML 使用 cloudpickle，以便使用 cuml.accel 训练的模型可以被加载并在 scikit-learn 中使用。

仅从可信来源反序列化模型。 pickle 模块（以及其扩展 joblib）并不安全：恶意载荷可以在反序列化过程中执行任意代码并危害你的系统。请勿从不可信或已被篡改的来源 unpickle 或加载数据。 这包括 pickle.load() / pickle.loads()、joblib.load() 以及任何基于文件的模型加载。详情与模式请参考 模型序列化与持久化 笔记本和 Python pickle 安全文档。

贡献

请参阅我们的 贡献指南。

参考资料

RAPIDS 团队发布了一系列包含深入技术解析和示例的博客文章。你可以在 Medium 上找到它们。

有关 cuML 背后技术的更多细节，以及 Python 机器学习领域的更广泛概述，请参阅 Machine Learning in Python: Main developments and technology trends in data science, machine learning, and artificial intelligence (2020)，作者 Sebastian Raschka, Joshua Patterson 和 Corey Nolet。

在项目中使用 cuML 时，请考虑引用该论文。你可以使用以下 BibTeX 引用：

@article{raschka2020machine,
  title={Machine Learning in Python: Main developments and technology trends in data science, machine learning, and artificial intelligence},
  author={Raschka, Sebastian and Patterson, Joshua and Nolet, Corey},
  journal={arXiv preprint arXiv:2002.04803},
  year={2020}
}

联系方式

更多详情请访问 RAPIDS 网站

开放式 GPU 数据科学

RAPIDS 开源软件库套件旨在使端到端的数据科学和分析流程完全在 GPU 上执行。它依赖 NVIDIA® CUDA® 原语进行底层计算优化，并通过用户友好的 Python 接口暴露 GPU 的并行性和高带宽内存速度。