注意
点击 此处 下载完整的示例代码
雅可比矩阵、Hessian 矩阵、hvp、vhp 等:函数变换的组合¶
计算雅可比矩阵或 Hessian 矩阵在许多非传统深度学习模型中很有用。使用 PyTorch 的常规自动微分 API(Tensor.backward()
、torch.autograd.grad
)来有效地计算这些量很困难(或很麻烦)。PyTorch 的 受 JAX 启发 的 函数变换 API 提供了有效计算各种高阶自动微分量的方法。
注意
本教程需要 PyTorch 2.0.0 或更高版本。
计算雅可比矩阵¶
import torch
import torch.nn.functional as F
from functools import partial
_ = torch.manual_seed(0)
让我们从一个想要计算雅可比矩阵的函数开始。这是一个具有非线性激活的简单线性函数。
让我们添加一些虚拟数据:权重、偏差和特征向量 x。
D = 16
weight = torch.randn(D, D)
bias = torch.randn(D)
x = torch.randn(D) # feature vector
让我们将 predict
视为一个将输入 x
从 \(R^D \to R^D\) 映射的函数。PyTorch Autograd 计算向量-雅可比矩阵积。为了计算此 \(R^D \to R^D\) 函数的完整雅可比矩阵,我们将不得不每次使用不同的单位向量逐行计算它。
def compute_jac(xp):
jacobian_rows = [torch.autograd.grad(predict(weight, bias, xp), xp, vec)[0]
for vec in unit_vectors]
return torch.stack(jacobian_rows)
xp = x.clone().requires_grad_()
unit_vectors = torch.eye(D)
jacobian = compute_jac(xp)
print(jacobian.shape)
print(jacobian[0]) # show first row
torch.Size([16, 16])
tensor([-0.5956, -0.6096, -0.1326, -0.2295, 0.4490, 0.3661, -0.1672, -1.1190,
0.1705, -0.6683, 0.1851, 0.1630, 0.0634, 0.6547, 0.5908, -0.1308])
与其逐行计算雅可比矩阵,我们可以使用 PyTorch 的 torch.vmap
函数变换来消除 for 循环并向量化计算。我们不能直接将 vmap
应用于 torch.autograd.grad
;相反,PyTorch 提供了一个 torch.func.vjp
变换,它与 torch.vmap
组合。
from torch.func import vmap, vjp
_, vjp_fn = vjp(partial(predict, weight, bias), x)
ft_jacobian, = vmap(vjp_fn)(unit_vectors)
# let's confirm both methods compute the same result
assert torch.allclose(ft_jacobian, jacobian)
在后续教程中,反向模式自动微分 (AD) 和 vmap
的组合将为我们提供每个样本的梯度。在本教程中,反向模式 AD 和 vmap
的组合为我们提供了雅可比矩阵计算!vmap
和自动微分变换的各种组合可以为我们提供不同的有趣量。
PyTorch 提供了 torch.func.jacrev
作为便利函数,它执行 vmap-vjp
组合来计算雅可比矩阵。jacrev
接受一个 argnums
参数,该参数表示我们希望针对哪个参数计算雅可比矩阵。
from torch.func import jacrev
ft_jacobian = jacrev(predict, argnums=2)(weight, bias, x)
# Confirm by running the following:
assert torch.allclose(ft_jacobian, jacobian)
让我们比较两种计算雅可比矩阵的方法的性能。函数转换版本的速度快得多(并且随着输出数量的增加而变得更快)。
通常,我们预计通过 vmap
进行向量化可以帮助消除开销并更好地利用您的硬件。
vmap
通过将外循环下推到函数的原始操作中来实现此魔力,从而获得更好的性能。
让我们快速创建一个函数来评估性能并处理微秒和毫秒测量。
def get_perf(first, first_descriptor, second, second_descriptor):
"""takes torch.benchmark objects and compares delta of second vs first."""
faster = second.times[0]
slower = first.times[0]
gain = (slower-faster)/slower
if gain < 0: gain *=-1
final_gain = gain*100
print(f" Performance delta: {final_gain:.4f} percent improvement with {second_descriptor} ")
然后运行性能比较。
from torch.utils.benchmark import Timer
without_vmap = Timer(stmt="compute_jac(xp)", globals=globals())
with_vmap = Timer(stmt="jacrev(predict, argnums=2)(weight, bias, x)", globals=globals())
no_vmap_timer = without_vmap.timeit(500)
with_vmap_timer = with_vmap.timeit(500)
print(no_vmap_timer)
print(with_vmap_timer)
<torch.utils.benchmark.utils.common.Measurement object at 0x7f41b8276ec0>
compute_jac(xp)
2.72 ms
1 measurement, 500 runs , 1 thread
<torch.utils.benchmark.utils.common.Measurement object at 0x7f41b09ecbb0>
jacrev(predict, argnums=2)(weight, bias, x)
691.61 us
1 measurement, 500 runs , 1 thread
让我们使用我们的 get_perf
函数对上述内容进行相对性能比较。
get_perf(no_vmap_timer, "without vmap", with_vmap_timer, "vmap")
Performance delta: 74.6171 percent improvement with vmap
此外,可以很容易地将问题反过来,说我们想计算模型参数(权重、偏差)而不是输入的雅可比矩阵。
# note the change in input via ``argnums`` parameters of 0,1 to map to weight and bias
ft_jac_weight, ft_jac_bias = jacrev(predict, argnums=(0, 1))(weight, bias, x)
反向模式雅可比 (jacrev
) 与正向模式雅可比 (jacfwd
)¶
我们提供了两个 API 来计算雅可比矩阵:jacrev
和 jacfwd
。
jacrev
使用反向模式 AD。如上所述,它是我们的vjp
和vmap
变换的组合。jacfwd
使用正向模式 AD。它实现为我们的jvp
和vmap
变换的组合。
jacfwd
和 jacrev
可以互相替换,但它们具有不同的性能特征。
作为一个一般的经验法则,如果您正在计算 \(R^N \to R^M\) 函数的雅可比矩阵,并且输出的数量远远大于输入的数量(例如,\(M > N\)),那么 jacfwd
是首选,否则使用 jacrev
。此规则有一些例外情况,但以下是一个非严格的论证。
在反向模式 AD 中,我们逐行计算雅可比矩阵,而在正向模式 AD(计算雅可比向量积)中,我们逐列计算它。雅可比矩阵有 M 行和 N 列,因此如果它在某一方面更高或更宽,我们可能更喜欢处理较少行或列的方法。
首先,让我们用更多输入比输出进行基准测试。
Din = 32
Dout = 2048
weight = torch.randn(Dout, Din)
bias = torch.randn(Dout)
x = torch.randn(Din)
# remember the general rule about taller vs wider... here we have a taller matrix:
print(weight.shape)
using_fwd = Timer(stmt="jacfwd(predict, argnums=2)(weight, bias, x)", globals=globals())
using_bwd = Timer(stmt="jacrev(predict, argnums=2)(weight, bias, x)", globals=globals())
jacfwd_timing = using_fwd.timeit(500)
jacrev_timing = using_bwd.timeit(500)
print(f'jacfwd time: {jacfwd_timing}')
print(f'jacrev time: {jacrev_timing}')
torch.Size([2048, 32])
jacfwd time: <torch.utils.benchmark.utils.common.Measurement object at 0x7f41386f4070>
jacfwd(predict, argnums=2)(weight, bias, x)
1.25 ms
1 measurement, 500 runs , 1 thread
jacrev time: <torch.utils.benchmark.utils.common.Measurement object at 0x7f41b08b1e40>
jacrev(predict, argnums=2)(weight, bias, x)
11.57 ms
1 measurement, 500 runs , 1 thread
然后进行相对基准测试。
get_perf(jacfwd_timing, "jacfwd", jacrev_timing, "jacrev", );
Performance delta: 825.6340 percent improvement with jacrev
现在反过来 - 输出 (M) 比输入 (N) 多。
Din = 2048
Dout = 32
weight = torch.randn(Dout, Din)
bias = torch.randn(Dout)
x = torch.randn(Din)
using_fwd = Timer(stmt="jacfwd(predict, argnums=2)(weight, bias, x)", globals=globals())
using_bwd = Timer(stmt="jacrev(predict, argnums=2)(weight, bias, x)", globals=globals())
jacfwd_timing = using_fwd.timeit(500)
jacrev_timing = using_bwd.timeit(500)
print(f'jacfwd time: {jacfwd_timing}')
print(f'jacrev time: {jacrev_timing}')
jacfwd time: <torch.utils.benchmark.utils.common.Measurement object at 0x7f41cfede2c0>
jacfwd(predict, argnums=2)(weight, bias, x)
6.29 ms
1 measurement, 500 runs , 1 thread
jacrev time: <torch.utils.benchmark.utils.common.Measurement object at 0x7f41b82772e0>
jacrev(predict, argnums=2)(weight, bias, x)
813.83 us
1 measurement, 500 runs , 1 thread
以及相对性能比较。
get_perf(jacrev_timing, "jacrev", jacfwd_timing, "jacfwd")
Performance delta: 672.6125 percent improvement with jacfwd
使用 functorch.hessian 计算 Hessian¶
我们提供了一个便利的 API 来计算 Hessian:torch.func.hessiani
。Hessian 是雅可比矩阵的雅可比矩阵(或偏导数的偏导数,即二阶导数)。
这表明人们只需组合 functorch 雅可比变换即可计算 Hessian。事实上,在幕后,hessian(f)
仅仅是 jacfwd(jacrev(f))
。
注意:为了提高性能:根据您的模型,您可能还需要使用 jacfwd(jacfwd(f))
或 jacrev(jacrev(f))
来计算 Hessian,利用上面关于更宽与更高矩阵的经验法则。
from torch.func import hessian
# lets reduce the size in order not to overwhelm Colab. Hessians require
# significant memory:
Din = 512
Dout = 32
weight = torch.randn(Dout, Din)
bias = torch.randn(Dout)
x = torch.randn(Din)
hess_api = hessian(predict, argnums=2)(weight, bias, x)
hess_fwdfwd = jacfwd(jacfwd(predict, argnums=2), argnums=2)(weight, bias, x)
hess_revrev = jacrev(jacrev(predict, argnums=2), argnums=2)(weight, bias, x)
让我们验证无论使用 Hessian API 还是使用 jacfwd(jacfwd())
,我们是否都得到了相同的结果。
True
批量雅可比矩阵和批量 Hessian¶
在以上示例中,我们一直在使用单个特征向量。在某些情况下,您可能希望针对一批输入获取一批输出的雅可比矩阵。也就是说,给定一个形状为 (B, N)
的输入批次和一个从 \(R^N \to R^M\) 的函数,我们希望得到一个形状为 (B, M, N)
的雅可比矩阵。
最简单的方法是使用 vmap
。
batch_size = 64
Din = 31
Dout = 33
weight = torch.randn(Dout, Din)
print(f"weight shape = {weight.shape}")
bias = torch.randn(Dout)
x = torch.randn(batch_size, Din)
compute_batch_jacobian = vmap(jacrev(predict, argnums=2), in_dims=(None, None, 0))
batch_jacobian0 = compute_batch_jacobian(weight, bias, x)
weight shape = torch.Size([33, 31])
如果您有一个从 (B, N) -> (B, M) 的函数,并且确定每个输入都会产生一个独立的输出,那么有时也可以在不使用 vmap
的情况下通过对输出求和然后计算该函数的雅可比矩阵来实现。
def predict_with_output_summed(weight, bias, x):
return predict(weight, bias, x).sum(0)
batch_jacobian1 = jacrev(predict_with_output_summed, argnums=2)(weight, bias, x).movedim(1, 0)
assert torch.allclose(batch_jacobian0, batch_jacobian1)
如果您有一个从 \(R^N \to R^M\) 的函数,但输入是批处理的,则将 vmap
与 jacrev
组合以计算批量雅可比矩阵。
最后,批量 Hessian 可以类似地计算。最容易想到的是使用 vmap
对 Hessian 计算进行批处理,但在某些情况下,求和技巧也有效。
compute_batch_hessian = vmap(hessian(predict, argnums=2), in_dims=(None, None, 0))
batch_hess = compute_batch_hessian(weight, bias, x)
batch_hess.shape
torch.Size([64, 33, 31, 31])
计算 Hessian 向量积¶
计算 Hessian 向量积 (hvp) 的朴素方法是具体化完整的 Hessian 并与向量执行点积。我们可以做得更好:事实证明,我们不需要具体化完整的 Hessian 就可以做到这一点。我们将介绍两种(在许多种中)计算 Hessian 向量积的不同策略:- 将反向模式 AD 与反向模式 AD 组合- 将反向模式 AD 与正向模式 AD 组合。
将反向模式 AD 与正向模式 AD 组合(而不是反向模式与反向模式)通常是计算 hvp 更节省内存的方式,因为正向模式 AD 不需要构建 Autograd 图并保存反向的中间结果。
以下是一些示例用法。
def f(x):
return x.sin().sum()
x = torch.randn(2048)
tangent = torch.randn(2048)
result = hvp(f, (x,), (tangent,))
如果 PyTorch 正向 AD 无法覆盖您的操作,那么我们可以改为将反向模式 AD 与反向模式 AD 组合。
脚本的总运行时间:(0 分 12.873 秒)