关于cuML UMAP在A100 GPU上非欧几里得度量支持情况的问询

我尝试用A100 GPU加速UMAP计算，发现欧几里得度量完全不适用，correlation、cosine这类度量效果更好，但GPU版本代码只能正常运行基于欧几里得度量的计算，CPU版本则无此问题。

使用工具版本 #

cuml                      23.04.01        cuda11_py310_230421_g958186d07_0    rapidsai
libcuml                   23.04.01        cuda11_230421_g958186d07_0          rapidsai
libcumlprims              23.04.00        cuda11_230412_g7502d8e_0            nvidia
python                    3.10.11         he550d4f_0_cpython                  conda-forge

相关代码 #

def umap_cpu(ip_mat, n_components, n_neighbors, metric):
    import umap
    from sklearn.preprocessing import StandardScaler

    scaler = StandardScaler()
    ip_std = scaler.fit_transform(ip_mat)

    reducer = umap.UMAP(n_components=n_components, n_neighbors=n_neighbors, metric=metric)
    umap_embed = reducer.fit_transform(ip_std)

    return umap_embed

def umap_gpu(ip_mat, n_components, n_neighbors, metric):
    import cuml
    from cuml.manifold import UMAP
    from sklearn.preprocessing import StandardScaler
    import os

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    scaler = StandardScaler()
    ip_std = scaler.fit_transform(ip_mat)

    reducer = UMAP(n_components=n_components, n_neighbors=n_neighbors, metric=metric)
    umap_embed = reducer.fit_transform(ip_std)

    return umap_embed

发现的矛盾信息 #

通过help命令查看工具文档，显示支持多种度量：

metric：字符串（默认值='euclidean'）
使用的距离度量。支持的距离包括['l1', 'cityblock', 'taxicab', 'manhattan', 'euclidean', 'l2', 'sqeuclidean', 'canberra', 'minkowski', 'chebyshev', 'linf', 'cosine', 'correlation', 'hellinger', 'hamming', 'jaccard']。需要参数的度量（如minkowski）可通过metric_kwds字典传递参数。

但找到一篇GitHub旧讨论提到：

PR 将允许修改输入KNN图的度量，但目前仅支持分类和欧几里得作为目标度量。我们可以支持更多目标度量（已有相关issue），但这需要在SGD中使用不同的目标函数。我认为当前Python异常抛出存在错误（本issue中已指出）

我想确认：目前是否已实现对其他度量的完整支持，还是help工具显示的信息有误？

解答 #

针对你使用的cuML 23.04版本，结论是help文档存在表述不准确，实际GPU版UMAP并未完全支持所有列出的度量：

• 你看到的GitHub讨论内容在23.04版本中仍然有效：当前GPU版UMAP仅支持欧几里得和分类作为目标度量（即降维优化阶段使用的度量）。
• help文档中列出的多种度量，仅适用于KNN图构建阶段，而非UMAP整个流程。也就是说，虽然可以用这些度量计算邻域，但后续的降维优化过程依然只能用欧几里得或分类度量，这会导致实际效果和CPU版UMAP（全程使用指定度量）不一致，甚至直接报错。
• 对于你需要的correlation、cosine这类度量，在cuML 23.04版本的GPU UMAP中无法完整支持，只能继续使用CPU版UMAP来获得预期效果。

若需GPU加速支持更多目标度量，可尝试升级到更高版本的cuML（如23.10及以后），后续版本逐步完善了非欧几里得度量的支持，建议升级前查阅对应版本的官方文档确认具体支持情况。

内容的提问来源于stack exchange，提问作者Quiescent