torch.func 快速入门

创建日期：2025 年 6 月 12 日 | 最后更新日期：2025 年 6 月 12 日

什么是 torch.func？

torch.func，之前称为 functorch，是一个在 PyTorch 中实现的类似 JAX 的可组合函数变换的库。

• “函数变换”是一种高阶函数，它接受一个数值函数并返回一个计算不同量的新函数。

• torch.func 具有自动微分变换（grad(f) 返回一个计算 f 的梯度的函数），向量化/批处理变换（vmap(f) 返回一个在输入批次上计算 f 的函数），以及其他功能。

• 这些函数变换可以任意组合。例如，组合 vmap(grad(f)) 可以计算出一个称为“每样本梯度”的量，这是当前标准 PyTorch 难以高效计算的。

为什么需要可组合函数变换？

目前在 PyTorch 中有一些难以实现的使用案例

• 计算逐样本梯度（或其他逐样本量）

• 在单机上运行模型集成

• 高效地批处理 MAML 内部循环中的任务

• 高效地计算雅可比矩阵和 Hessian 矩阵

• 高效地计算批处理雅可比矩阵和 Hessian 矩阵

组合 vmap()、grad()、vjp() 和 jvp() 变换，使我们能够通过为每个变换设计一个单独的子系统来表达上述功能。

有哪些变换？

grad() (梯度计算)

grad(func) 是我们的梯度计算变换。它返回一个新的函数，该函数计算 func 的梯度。它假设 func 返回一个单一元素的 Tensor，并且默认情况下它计算 func 的输出相对于第一个输入的梯度。

import torch
from torch.func import grad
x = torch.randn([])
cos_x = grad(lambda x: torch.sin(x))(x)
assert torch.allclose(cos_x, x.cos())

# Second-order gradients
neg_sin_x = grad(grad(lambda x: torch.sin(x)))(x)
assert torch.allclose(neg_sin_x, -x.sin())

vmap() (自动向量化)

注意：vmap() 对其可以使用的代码施加了限制。有关更多详细信息，请参阅 UX 限制。

vmap(func)(*inputs) 是一种变换，它为 func 中的所有 Tensor 操作添加了一个维度。vmap(func) 返回一个新函数，该函数将 func 映射到 inputs 中每个 Tensor 的某个维度（默认为 0）。

vmap 对于隐藏批次维度非常有用：您可以编写一个处理单个样本的函数 func，然后使用 vmap(func) 将其提升为可以处理样本批次的函数，从而获得更简单的建模体验。

import torch
from torch.func import vmap
batch_size, feature_size = 3, 5
weights = torch.randn(feature_size, requires_grad=True)

def model(feature_vec):
    # Very simple linear model with activation
    assert feature_vec.dim() == 1
    return feature_vec.dot(weights).relu()

examples = torch.randn(batch_size, feature_size)
result = vmap(model)(examples)

当与 grad() 组合时，vmap() 可用于计算每样本梯度。

from torch.func import vmap
batch_size, feature_size = 3, 5

def model(weights,feature_vec):
    # Very simple linear model with activation
    assert feature_vec.dim() == 1
    return feature_vec.dot(weights).relu()

def compute_loss(weights, example, target):
    y = model(weights, example)
    return ((y - target) ** 2).mean()  # MSELoss

weights = torch.randn(feature_size, requires_grad=True)
examples = torch.randn(batch_size, feature_size)
targets = torch.randn(batch_size)
inputs = (weights,examples, targets)
grad_weight_per_example = vmap(grad(compute_loss), in_dims=(None, 0, 0))(*inputs)

vjp() (向量-雅可比乘积)

该 vjp() 变换将 func 应用于 inputs，并返回一个新的函数，该函数在给定一些 cotangents Tensor 后计算向量-雅可比乘积 (vjp)。

from torch.func import vjp

inputs = torch.randn(3)
func = torch.sin
cotangents = (torch.randn(3),)

outputs, vjp_fn = vjp(func, inputs); vjps = vjp_fn(*cotangents)

jvp() (雅可比-向量乘积)

该 jvp() 变换计算雅可比-向量乘积，也称为“前向模式 AD”。它不像大多数其他变换那样是高阶函数，但它会返回 func(inputs) 的输出以及 jvp。

from torch.func import jvp
x = torch.randn(5)
y = torch.randn(5)
f = lambda x, y: (x * y)
_, out_tangent = jvp(f, (x, y), (torch.ones(5), torch.ones(5)))
assert torch.allclose(out_tangent, x + y)

jacrev()、jacfwd() 和 hessian()

该 jacrev() 变换返回一个新函数，该函数接收 x 并使用反向模式 AD 返回函数相对于 x 的雅可比矩阵。

from torch.func import jacrev
x = torch.randn(5)
jacobian = jacrev(torch.sin)(x)
expected = torch.diag(torch.cos(x))
assert torch.allclose(jacobian, expected)

可以将 jacrev() 与 vmap() 组合以生成批处理雅可比矩阵。

x = torch.randn(64, 5)
jacobian = vmap(jacrev(torch.sin))(x)
assert jacobian.shape == (64, 5, 5)

jacfwd() 是 jacrev 的即插即用替代品，它使用前向模式 AD 计算雅可比矩阵。

from torch.func import jacfwd
x = torch.randn(5)
jacobian = jacfwd(torch.sin)(x)
expected = torch.diag(torch.cos(x))
assert torch.allclose(jacobian, expected)

将 jacrev() 与自身或 jacfwd() 组合可以生成 Hessian 矩阵。

def f(x):
    return x.sin().sum()

x = torch.randn(5)
hessian0 = jacrev(jacrev(f))(x)
hessian1 = jacfwd(jacrev(f))(x)

hessian() 是一个方便函数，它结合了 jacfwd 和 jacrev。

from torch.func import hessian

def f(x):
    return x.sin().sum()

x = torch.randn(5)
hess = hessian(f)(x)