快捷方式

torch.func API 参考

函数变换

vmap

vmap 是向量化映射;vmap(func) 返回一个新函数,该函数在输入的某些维度上映射 func

grad

grad 运算符有助于计算 func 相对于由 argnums 指定的输入的梯度。

grad_and_value

返回一个函数,用于计算梯度和原始(或前向)计算的元组。

vjp

代表向量-雅可比积,返回一个元组,其中包含应用于 primalsfunc 的结果以及一个函数,当给定 cotangents 时,该函数计算 func 相对于 primals 的反向模式雅可比乘以 cotangents

jvp

代表雅可比-向量积,返回一个元组,其中包含 func(*primals) 的输出以及“在 primals 处计算的 func 的雅可比”乘以 tangents

linearize

返回 funcprimals 处的取值以及在 primals 处的线性近似。

jacrev

使用反向模式自动微分计算 func 相对于索引为 argnum 的参数的雅可比。

jacfwd

使用前向模式自动微分计算 func 相对于索引为 argnum 的参数的雅可比。

hessian

通过前向-反向策略计算 func 相对于索引为 argnum 的参数的 Hessian。

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)

但是,如果您想执行类似于计算模型参数上的雅可比的操作,则需要一种方法来构建一个函数,其中参数是函数的输入。这就是 functional_call() 的作用:它接受一个 nn.Module、变换后的 parameters 以及 Module 的前向传递的输入。它返回使用替换的参数运行 Module 的前向传递的值。

以下是如何计算参数上的雅可比的示例

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

def f(params, x):
    return torch.func.functional_call(model, params, x)

x = torch.randn(3)
jacobian = jacrev(f)(dict(model.named_parameters()), x)

functional_call

通过使用提供的参数替换模块参数和缓冲区来对模块执行函数调用。

stack_module_state

准备一个 torch.nn.Modules 列表,以便与 vmap() 一起进行集成。

replace_all_batch_norm_modules_

就地更新 root,方法是将 running_meanrunning_var 设置为 None,并将 root 中的任何 nn.BatchNorm 模块的 track_running_stats 设置为 False。

如果您正在寻找有关修复 Batch Norm 模块的信息,请按照此处的指南进行操作

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源