跳转到主要内容
博客

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

线性代数是深度学习和科学计算的核心，它一直是 PyTorch 的核心组成部分。PyTorch 1.9 通过 torch.linalg 模块扩展了 PyTorch 对线性代数运算的支持。此模块（此处有文档记录：https://pytorch.ac.cn/docs/master/linalg.html?highlight=linalg#module-torch.linalg）包含 26 个运算符，其中包括更快、更易于使用的旧版 PyTorch 运算符，以及 NumPy 线性代数模块的每个函数（https://numpy.com.cn/doc/stable/reference/routines.linalg.html）的扩展，支持加速器和 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 的线性代数模块使用 LAPACK 标准的 OpenBLAS 或 MKL 实现来最大化其 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 的新线性代数功能！