注意

RAFT 中的向量搜索和聚类算法正在迁移到一个新的专门用于向量搜索的库 cuVS。在此迁移过程中,我们将继续支持 RAFT 中的向量搜索算法,但在 RAPIDS 24.06(6月)版本后将不再更新。我们计划在 RAPIDS 24.10(10月)版本前完成迁移,并在 24.12(12月)版本中将其从 RAFT 中完全移除。

随机#

此页面提供了 pylibraft.random 包公开元素的 pylibraft 类参考。

pylibraft.random.rmat(out, theta, r_scale, c_scale, seed=12345, handle=None)[source]#

根据输入分布生成 RMAT 邻接列表。

参数:
out: 符合 CUDA 数组接口的矩阵,形状为 (n_edges, 2)。这其中包含按对连续存储的源/目标节点 ID。

包含按对连续存储的 src/dst 节点 ID。

theta: 符合 CUDA 数组接口的矩阵,形状为

(max(r_scale, c_scale) * 4) 这存储了每个 RMAT 级别的概率分布

r_scale: 源节点数的 log2
c_scale: 目标节点数的 log2
seed: 用于重现性的随机种子
handle可选的 RAFT 资源句柄,用于重用 CUDA 资源。

如果未提供句柄,CUDA 资源将在函数内部分配,并在函数退出前同步。如果提供了句柄,您需要在访问输出之前通过调用 handle.sync() 来明确进行同步。

示例

>>> import cupy as cp
>>> from pylibraft.common import Handle
>>> from pylibraft.random import rmat
>>> n_edges = 5000
>>> r_scale = 16
>>> c_scale = 14
>>> theta_len = max(r_scale, c_scale) * 4
>>> out = cp.empty((n_edges, 2), dtype=cp.int32)
>>> theta = cp.random.random_sample(theta_len, dtype=cp.float32)
>>> # A single RAFT handle can optionally be reused across
>>> # pylibraft functions.
>>> handle = Handle()
>>> rmat(out, theta, r_scale, c_scale, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()