跳转到主要内容
博客

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

线性代数对于深度学习和科学计算至关重要，它一直是 PyTorch 的核心组成部分。PyTorch 1.9 通过 torch.linalg 模块扩展了 PyTorch 对线性代数运算的支持。这个模块此处有文档，包含 26 个运算符，其中包括更快速、更易于使用的旧 PyTorch 运算符版本，NumPy 线性代数模块中的所有函数都扩展了加速器和 Autograd 支持，还有一些完全新的运算符。这使得 torch.linalg 对 NumPy 用户来说立即熟悉，也是 PyTorch 线性代数支持的一个令人兴奋的更新。

PyTorch 中类似 NumPy 的线性代数

如果您熟悉 NumPy 的线性代数模块，那么开始使用 torch.linalg 将会很容易。在大多数情况下，它可以直接替换。让我们以使用Cholesky 分解从多元正态分布中抽取样本为例，演示这一点。

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 默认打印张量的方式。

Cholesky 分解还可以帮助我们快速计算非退化多元正态分布的概率密度函数。该计算中昂贵的项之一是协方差矩阵行列式的平方根。然而，利用行列式的性质和 Cholesky 分解，我们可以比朴素计算更快地计算出相同的结果。以下是演示这一点的 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

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

Autograd 支持

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

让我们看一个非常简单的程序，它只计算一个逆矩阵和该操作的梯度，以展示 Autograd 的工作原理。

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]])

我们可以通过自己定义 autograd 公式来模仿 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]]

当然，随着程序变得更加复杂，内置的 Autograd 支持会更加方便，PyTorch 的线性代数模块支持实数和复数的 Autograd。

CUDA 支持

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

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

这些图表显示，在大矩阵上的性能显著提升，并且批处理性能全面更好。其他线性代数操作，包括 torch.linalg.qr 和 torch.linalg.lstsq，也改进了其 CUDA 性能。

超越 NumPy

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

PyTorch 线性代数的未来

torch.linalg 模块速度快且熟悉，并对 Autograd 和加速器提供了很好的支持。它也已经被 botorch 等库使用。但我们不止于此。我们计划继续更新更多 PyTorch 现有线性代数功能（如 torch.lobpcg），并为低秩和稀疏线性代数提供更多支持。我们也希望听到您关于如何改进的反馈，因此请在论坛上开始对话，或在我们的Github 上提交问题并分享您的想法。

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