博客

torch.linalg 模块：PyTorch 中带有自动微分的加速线性代数

线性代数是深度学习和科学计算的基础，也是 PyTorch 的核心组成部分。PyTorch 1.9 通过 torch.linalg 模块扩展了对线性代数运算的支持。该模块在此处有相关文档，包含 26 个算子，其中包括比旧版 PyTorch 算子更快速、更易用的版本，涵盖了 NumPy 线性代数模块中的所有函数（并增加了加速器和自动微分支持），以及一些全新的算子。这使得 torch.linalg 对于 NumPy 用户来说立即可用，也是 PyTorch 线性代数支持方面令人振奋的更新。

PyTorch 中类 NumPy 的线性代数

如果您熟悉 NumPy 的线性代数模块，那么上手 torch.linalg 将非常容易。在大多数情况下，它是一个直接替代方案。让我们以使用乔列斯基分解（Cholesky decomposition）从多元正态分布中抽取样本为例来演示这一点。

import numpy as np

# Creates inputs
np.random.seed(0)
mu_np = np.random.rand(4)
L = np.random.rand(4, 4)
# Covariance matrix sigma is positive-definite
sigma_np = L @ L.T + np.eye(4)
normal_noise_np = np.random.standard_normal(mu_np.size)

def multivariate_normal_sample_np(mu, sigma, normal_noise):
    return mu + np.linalg.cholesky(sigma) @ normal_noise

print("Random sample: ", 
      multivariate_normal_sample_np(mu_np, sigma_np, normal_noise_np))
: Random sample: [2.9502426 1.78518077 1.83168697 0.90798228]

现在让我们看看在 PyTorch 中实现的相同采样器。

import torch

def multivariate_normal_sample_torch(mu, sigma, normal_noise):
    return mu + torch.linalg.cholesky(sigma) @ normal_noise

这两个函数是完全相同的，我们可以通过传入封装为 PyTorch 张量的相同参数来验证它们的行为。

# NumPy arrays are wrapped as tensors and share their memory
mu_torch = torch.from_numpy(mu_np)
sigma_torch = torch.from_numpy(sigma_np)
normal_noise_torch = torch.from_numpy(normal_noise_np)

multivariate_normal_sample_torch(mu_torch, sigma_torch, normal_noise_torch)
: tensor([2.9502, 1.7852, 1.8317, 0.9080], dtype=torch.float64)

唯一的区别在于 PyTorch 默认打印张量的方式。

乔列斯基分解还可以帮助我们快速计算非退化多元正态分布的概率密度函数。该计算中代价昂贵的一项是协方差矩阵行列式的平方根。使用行列式的性质和乔列斯基分解，我们可以比普通计算更快地得到结果。以下是演示这一点的 NumPy 程序：

sqrt_sigma_det_np = np.sqrt(np.linalg.det(sigma_np))
sqrt_L_det_np = np.prod(np.diag(np.linalg.cholesky(sigma_np)))

print("|sigma|^0.5 = ", sqrt_sigma_det_np)
: |sigma|^0.5 = 4.237127491242027
 
print("|L| = ", sqrt_L_det_np)
: |L| = 4.237127491242028

这是在 PyTorch 中进行的相同验证：

sqrt_sigma_det_torch = torch.sqrt(torch.linalg.det(sigma_torch))
sqrt_L_det_torch = torch.prod(torch.diag(torch.linalg.cholesky(sigma_torch)))

print("|sigma|^0.5 = ", sqrt_sigma_det_torch)
: |sigma|^0.5 = tensor(4.2371, dtype=torch.float64) 

print("|L| = ", sqrt_L_det_torch)
: |L| = tensor(4.2371, dtype=torch.float64)

我们可以使用 PyTorch 内置的基准测试工具来测量运行时间的差异。

import torch.utils.benchmark as benchmark

t0 = benchmark.Timer(
    stmt='torch.sqrt(torch.linalg.det(sigma))',
    globals={'sigma': sigma_torch})

t1 = benchmark.Timer(
    stmt='torch.prod(torch.diag(torch.linalg.cholesky(sigma)))',
    globals={'sigma': sigma_torch})

print(t0.timeit(100))
: torch.sqrt(torch.linalg.det(sigma))
  80.80 us
  1 measurement, 100 runs , 1 thread


print(t1.timeit(100))
: torch.prod(torch.diag(torch.linalg.cholesky(sigma)))
  11.56 us
  1 measurement, 100 runs , 1 thread

这表明使用乔列斯基分解的方法可以显著提高速度。在幕后，PyTorch 的线性代数模块使用 OpenBLAS 或 MKL 实现的 LAPACK 标准，以最大限度地提高 CPU 性能。

自动微分支持

PyTorch 的线性代数模块不仅实现了与 NumPy 线性代数模块相同的功能（以及更多功能），还通过自动微分和 CUDA 支持对它们进行了扩展。

让我们看一个非常简单的程序，它只计算逆矩阵和该运算的梯度，以展示自动微分的工作原理：

t = torch.tensor(((1, 2), (3, 4)), dtype=torch.float32, requires_grad=True)

inv = torch.linalg.inv(t)
inv.backward(torch.ones_like(inv))

print(t.grad)
: tensor([[-0.5000, 0.5000],
          [ 0.5000, -0.5000]])

我们可以通过自己定义自动微分公式，在 NumPy 中模拟同样的计算：

a = np.array(((1, 2), (3, 4)), dtype=np.float32)

inv_np = np.linalg.inv(a)

def inv_backward(result, grad):
    return -(result.transpose(-2, -1) @ (grad @ result.transpose(-2, -1)))
grad_np = inv_backward(inv_np, np.ones_like(inv_np))

print(grad_np)
: [[-0.5 0.5]
   [ 0.5 -0.5]]

当然，随着程序变得越来越复杂，拥有内置的自动微分支持会非常方便，而且 PyTorch 的线性代数模块同时支持实数和复数的自动微分。

CUDA 支持

对自动微分和加速器（如 CUDA 设备）的支持是 PyTorch 的核心部分。torch.linalg 模块是与 NVIDIA 的 PyTorch 和 cuSOLVER 团队共同开发的，他们通过 cuSOLVER、cuBLAS 和 MAGMA 库帮助优化了其在 CUDA 设备上的性能。这些改进使 PyTorch 的 CUDA 线性代数运算比以往任何时候都更快。例如，让我们看看 PyTorch 1.9 的 torch.linalg.cholesky 与 PyTorch 1.8（现已弃用）的 torch.cholesky 的性能对比：

（以上图表使用 Ampere A100 GPU、CUDA 11.3、cuSOLVER 11.1.1.58 和 MAGMA 2.5.2 创建。矩阵采用双精度。）

这些图表显示，在较大的矩阵上性能得到了显著提升，且批处理性能在各方面都有所改善。其他线性代数运算，包括 torch.linalg.qr 和 torch.linalg.lstsq，其 CUDA 性能也得到了提升。

超越 NumPy

除了提供 NumPy 线性代数模块中所有函数并增加对自动微分和加速器的支持外，torch.linalg 还有一些自己的新函数。NumPy 的 linalg.norm 不允许用户在任意维度子集上计算向量范数，因此为了实现此功能，我们添加了 torch.linalg.vector_norm。我们也已经开始更新 PyTorch 中其他的线性代数功能，因此我们创建了 torch.linalg.householder_product 来替代旧的 torch.orgqr，并且我们计划在未来继续增加更多的线性代数功能。

PyTorch 线性代数的未来

torch.linalg 模块快速、易用，并对自动微分和加速器提供了极好的支持。它也已经在像 botorch 这样的库中得到应用。但我们不会就此止步。我们计划继续更新 PyTorch 现有的更多线性代数功能（如 torch.lobpcg），并提供对低秩和稀疏线性代数的更多支持。我们也希望听到您关于如何改进的反馈，欢迎在论坛上发起讨论，或者在我们的 Github 上提交问题，分享您的想法。

我们期待您的反馈，并期待看到社区如何利用 PyTorch 的新线性代数功能！