logo

请访问我们的网站获取文档、贡献指南和教程。

在 GPU 上进行内核运算，支持自动微分，无内存溢出

KeOps 库允许你计算大型数组的归约操作，其中数组的元素由数学公式或神经网络给出。它结合了高效的 C++ 例程和自动微分引擎，可与 Python（NumPy、PyTorch）、Matlab 和 R 一起使用。

它非常适合计算内核矩阵向量乘积、K 近邻查询、N 体相互作用、点云卷积及其相关梯度。关键是，即使相应的内核或距离矩阵不适合 RAM 或 GPU 内存，它也能表现出色。与 PyTorch GPU 基准相比，KeOps 在广泛的几何应用中提供了10-100 倍的加速，从内核方法到几何深度学习。

符号矩阵

为什么使用 KeOps？ 数学库将大多数对象表示为矩阵和张量：

(a) 密集矩阵。 变量通常被编码为密集数值数组 Mi,j = M[i,j]。这种表示方法便捷且得到广泛支持，但也给计算机内存带来了沉重负担。不幸的是，大型数组在移动时代价高昂，甚至可能无法装入 RAM 或 GPU 内存。

实际上，这意味着大多数科学程序都受到内存限制。大多数神经网络和数学计算的运行时间并非受限于 CPU 和 CUDA 核心的原始能力，而是受限于大型数组从内存电路到算术计算单元的耗时传输。
(b) 稀疏矩阵。 为解决这个问题，一个常见的解决方案是依赖稀疏矩阵：具有少量非零系数的张量。我们使用索引列表 (in,jn) 和值 Mn = Min,jn 来表示这些对象，它们对应于少量非零条目。然后使用索引方法和分散内存访问来实现矩阵向量运算。

这种方法很优雅，允许我们用较小的内存占用表示大型数组。但不幸的是，它在 GPU 上流式处理效果不佳：并行计算设备被设计为执行块式内存访问，难以处理随机索引 (in,jn) 列表。因此，与密集数组相比，稀疏编码只有在矩阵非零系数少于 1% 时才能加速计算。这种限制性条件使得稀疏矩阵在图形和网格处理之外的用途有限。
(c) 符号矩阵。 KeOps 提供了另一种加速张量程序的解决方案。我们的关键发现是，机器学习和应用数学中使用的大多数大型数组都具有共同的数学结构。距离矩阵、内核矩阵、点云卷积和注意力层都可以描述为符号张量：给定两个向量集合 (xi) 和 (yj)，它们在位置 (i,j) 处的系数 Mi,j 由数学公式 F(xi,yj) 给出，该公式在数据样本 xi 和 yj 上求值。

这些对象在传统意义上并不"稀疏"...但仍然可以使用数学公式 F 和相对较小的数据数组 (xi) 和 (yj) 高效地描述。KeOps 库的主要目的是为这种抽象提供支持，同时具备深度学习库的所有优点：
- 透明的接口，支持 CPU 和 GPU 集成。
- 大量教程和基准测试。
- 全面支持自动微分、批处理和近似计算。

实际上，KeOps 符号张量既快速又内存高效。我们利用 CUDA 寄存器的结构来绕过算术和内存电路之间昂贵的内存传输。这使我们能够在广泛的设置中为 PyTorch GPU 程序提供10-100 倍的加速。

使用我们的 Python 接口，一个典型的代码示例如下：

# 在 GPU 上创建两个具有 3 列和（巨大）行数的数组
import torch  # 也支持 NumPy、Matlab 和 R
M, N, D = 1000000, 2000000, 3
x = torch.randn(M, D, requires_grad=True).cuda()  # x.shape = (1e6, 3)
y = torch.randn(N, D).cuda()                      # y.shape = (2e6, 3)

# 将我们的密集张量转换为 KeOps 符号变量，在位置 0 和 1 处有"虚拟"维度（用于"i"和"j"索引）：
from pykeops.torch import LazyTensor
x_i = LazyTensor(x.view(M, 1, D))  # x_i.shape = (1e6, 1, 3)
y_j = LazyTensor(y.view(1, N, D))  # y_j.shape = ( 1, 2e6,3)

# 现在我们可以进行大规模计算，而不会出现内存溢出：
D_ij = ((x_i - y_j)**2).sum(dim=2)  # 符号 (1e6,2e6,1) 平方距离矩阵
K_ij = (- D_ij).exp()               # 符号 (1e6,2e6,1) 高斯核矩阵

# 我们通过在两个"符号"维度 0 和 1 之一上使用归约操作（如 .sum()、.logsumexp() 或 .argmin()）
# 回到普通的 PyTorch 张量或 NumPy 数组。
# 这里，核密度估计 a_i = sum_j exp(-|x_i-y_j|^2)
# 使用具有线性内存占用的 CUDA 方案计算，
# 性能比标准 PyTorch 实现提高两个数量级。
a_i = K_ij.sum(dim=1)  # 真正的 torch.cuda.FloatTensor，a_i.shape = (1e6, 1)，

# 关键是，KeOps 完全支持自动微分！
g_x = torch.autograd.grad((a_i ** 2).sum(), [x])

KeOps 让你能够充分利用硬件而不影响易用性。它提供：

众多计算类型的线性（而非二次）内存占用。
支持广泛的可自由组合的数学公式。
无缝计算导数和梯度，可达任意阶。
Sum、LogSumExp、Min、Max 以及 ArgMin、ArgMax 或 K-min 归约。
用于大规模样条插值和高斯过程回归的共轭梯度求解器。
与标准包的透明集成，如用于线性代数的 SciPy 求解器。
块稀疏和从粗到细策略的接口。
支持多 GPU 配置。

更多详细信息请参见：

使用 KeOps 的项目

符号矩阵之于几何学习，就像稀疏矩阵之于图处理。因此，KeOps 可以在广泛的领域中使用，从形状分析（配准、几何深度学习、最优传输...）到机器学习（核方法、K-means、UMAP...）、高斯过程、计算生物学和物理学。

KeOps 为以下项目和库提供核心例程：

GPyTorch（来自康奈尔大学、哥伦比亚大学、宾夕法尼亚大学）和 Falkon（来自热那亚大学和 Sierra Inria 团队），两个用于高斯过程回归的库，现在可以扩展到十亿级数据集。
Deformetrica，来自 Aramis Inria 团队的计算解剖学软件。
Gudhi 库，用于拓扑数据分析和高维几何理解，来自 DataShape Inria 团队。
GeomLoss，一个用于 Chamfer（Hausdorff）距离、Kernel（Sobolev）散度和地球移动（Wasserstein）距离的 PyTorch 包。它提供的最优传输求解器可以在几秒钟内处理数百万个样本。
深度图匹配共识模块，用于学习和精炼图之间的结构对应关系。
FshapesTk 和 Shapes 工具箱，两个面向研究的 LDDMM 工具包。
HyenaDNA，用于 S4D 内核中使用的范德蒙德矩阵乘法内核和归约的并行计算。

许可、引用、学术使用

本库基于宽松的 MIT 许可证发布，完全兼容学术和商业应用。

如果您在研究论文中使用此代码，请引用我们的原始出版物：

Charlier, B., Feydy, J., Glaunès, J. A., Collin, F.-D. & Durif, G. 在GPU上进行核操作，具有自动微分功能，且无内存溢出。机器学习研究期刊 22, 1–6 (2021)。

@article{JMLR:v22:20-275,
  author  = {Benjamin Charlier and Jean Feydy and Joan Alexis Glaunès and François-David Collin and Ghislain Durif},
  title   = {在GPU上进行核操作，具有自动微分功能，且无内存溢出},
  journal = {机器学习研究期刊},
  year    = {2021},
  volume  = {22},
  number  = {74},
  pages   = {1-6},
  url     = {http://jmlr.org/papers/v22/20-275.html}
}

关于几何（深度）学习的应用，您也可以参考我们的NeurIPS 2020论文：

@article{feydy2020fast,
    title={使用符号矩阵的快速几何学习},
    author={Feydy, Jean and Glaun{\`e}s, Joan and Charlier, Benjamin and Bronstein, Michael},
    journal={神经信息处理系统进展},
    volume={33},
    year={2020}
}