旋风之旅¶
什么是functorch？¶
functorch 是一个用于在 PyTorch 中实现类似 JAX 的可组合函数变换的库。
“函数变换”是一种高阶函数，它接受一个数值函数并返回一个计算不同量的新函数。
functorch 具有自动微分变换（
grad(f)返回一个计算f梯度的函数），向量化/批处理变换（vmap(f)返回一个在输入批次上计算f的函数），以及其他功能。这些函数变换可以任意组合。例如，组合
vmap(grad(f))计算一个称为每样本梯度的量，这是目前PyTorch无法高效计算的。
此外，我们还在functorch.compile命名空间中提供了一个实验性的编译转换。我们的编译转换名为AOT（提前）Autograd，它会返回一个FX图（可选地包含一个反向传递），通过各种后端进行编译是您可以采取的一种路径。
为什么选择可组合函数转换？¶
在PyTorch中，有许多用例目前难以实现：
计算每个样本的梯度（或其他每个样本的量）
在单台机器上运行模型集合
在MAML的内循环中高效地批量处理任务
高效计算雅可比矩阵和海森矩阵
高效计算批量雅可比矩阵和海森矩阵
组合 vmap、grad、vjp 和 jvp 转换使我们能够表达上述内容，而无需为每个内容设计单独的子系统。
什么是变换？¶
grad（梯度计算）¶
grad(func) 是我们的梯度计算变换。它返回一个新函数，用于计算 func 的梯度。它假设 func 返回一个单元素张量，并且默认情况下，它计算 func 输出相对于第一个输入的梯度。
import torch
from functorch 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 对它可以使用的代码施加了限制。有关更多详细信息，请阅读其文档字符串。
vmap(func)(*inputs) 是一种转换，它为 func 中的所有张量操作添加了一个维度。vmap(func) 返回一个新函数，该函数将 func 映射到输入中每个张量的某个维度（默认：0）。
vmap 对于隐藏批次维度非常有用：可以编写一个在示例上运行的函数 func，然后将其提升为可以使用 vmap(func) 处理批次示例的函数，从而简化建模体验：
import torch
from functorch 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 functorch 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 张量的向量-雅可比积（vjp）。
from functorch 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 变换计算雅可比-向量积，也被称为“前向模式自动微分”。与大多数其他变换不同，它不是高阶函数，但它返回 func(inputs) 的输出以及雅可比-向量积。
from functorch import jvp
x = torch.randn(5)
y = torch.randn(5)
f = lambda x, y: (x * y)
_, output = jvp(f, (x, y), (torch.ones(5), torch.ones(5)))
assert torch.allclose(output, x + y)
jacrev, jacfwd, 和 hessian¶
jacrev 转换返回一个新函数，该函数接受 x 并使用反向模式 AD 返回函数相对于 x 的雅可比矩阵。
from functorch 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 functorch import jacfwd
x = torch.randn(5)
jacobian = jacfwd(torch.sin)(x)
expected = torch.diag(torch.cos(x))
assert torch.allclose(jacobian, expected)
将 jacrev 与自身或 jacfwd 组合可以生成海森矩阵：
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 functorch import hessian
def f(x):
return x.sin().sum()
x = torch.randn(5)
hess = hessian(f)(x)
