functorch

警告

我们已将 functorch 集成到 PyTorch 中。作为集成的最后一步，从 PyTorch 2.0 开始，functorch API 已被弃用。请改用 torch.func API，并查看 迁移指南 和 文档 了解更多详细信息。

函数变换

vmap	vmap 是向量化映射；vmap(func) 返回一个新函数，该函数在输入的某些维度上映射 func。
grad	grad 运算符有助于计算 func 相对于由 argnums 指定的输入的梯度。
grad_and_value	返回一个函数来计算梯度和原始（或正向）计算的元组。
vjp	代表向量-雅可比矩阵积，返回一个元组，包含应用于 primals 的 func 的结果，以及一个函数，该函数在给定 cotangents 时，计算 func 相对于 primals 的反向模式雅可比矩阵乘以 cotangents。
jvp	代表雅可比矩阵-向量积，返回一个元组，包含 func(*primals) 的输出，以及“在 primals 处评估的 func 的雅可比矩阵”乘以 tangents。
jacrev	使用反向模式自动微分计算 func 相对于索引为 argnum 的参数的雅可比矩阵。
jacfwd	使用正向模式自动微分计算 func 相对于索引为 argnum 的参数的雅可比矩阵。
hessian	通过正向-反向策略计算 func 相对于索引为 argnum 的参数的海森矩阵。
functionalize	functionalize 是一种变换，可用于从函数中移除（中间）变异和别名，同时保留函数的语义。

用于处理 torch.nn.Modules 的实用程序

一般来说，你可以对调用 torch.nn.Module 的函数进行变换。例如，以下是计算接受三个值并返回三个值的函数的雅可比矩阵的示例

model = torch.nn.Linear(3, 3)

def f(x):
    return model(x)

x = torch.randn(3)
jacobian = jacrev(f)(x)
assert jacobian.shape == (3, 3)

但是，如果你想做一些类似于计算模型参数的雅可比矩阵的事情，那么需要有一种方法来构造一个函数，其中参数是函数的输入。这就是 make_functional() 和 make_functional_with_buffers() 的目的：给定一个 torch.nn.Module，它们会返回一个新函数，该函数接受 parameters 和 Module 前向传递的输入。

make_functional	给定一个 torch.nn.Module，make_functional() 会提取状态（参数）并返回模型的函数版本 func。
make_functional_with_buffers	make_functional(model, disable_autograd_tracking=False) -> func, params
combine_state_for_ensemble	为使用 vmap() 集成准备一个 torch.nn.Modules 列表。

如果你正在寻找有关修复 Batch Norm 模块的信息，请遵循此处提供的指南

• 修补 Batch Norm