PyTorch GPU 中的并行 Cholesky 分解

Question

要获得包含 D 维正定矩阵 C 的行列式的可微项（在我的例子中是多元高斯的微分熵），我可以使用：

torch.log2(torch.potrf(C).diag()).sum() + D / 2.0 * (np.log2(2 * np.pi * np.e))

potrf(C) 执行 Cholesky 分解，其对角线元素的对数值总和为对数行列式除以 2。

我想在小批量矩阵上调用 potrf，以便在形状为 (N, D, D) 的张量上调用 potrf 会产生 N 不同的 Cholesky 分解。

目前我只能在Python循环中重复调用potrf()，这是对GPU并行计算能力的一种不良利用，结果运行速度比[=慢3倍左右30=].

是否可以在 GPU 上与 PyTorch 并行启动 Cholesky 分解？

Answer 1

PyTorch 现在可以使用批量 Cholesky 分解。

连同 batch inverse() 等

对于旧版本的 Pytorch

您正在寻找 Batch Cholesky 分解。目前在Pytorch中还没有实现，但是有一个open issue，打算以后加上。

我只知道 Batch LU factorization 在 Pytorch 0.4 中可用。你可以用它来获得类似的东西：

det(D) = det(P)det(L)det(U)

其中P的行列式是(-1)^t，P和U的行列式是对角线元素的乘积。

Answer 2

从 1.8 版开始，PyTorch 原生支持 numpy 风格的 torch.linalg 操作，包括 Cholesky 分解：

torch.linalg.cholesky(input)

NOTE: When given inputs on a CUDA device, this function synchronizes that device with the CPU.

Parallel Cholesky decomposition in PyTorch GPU