快捷方式

torch.func

torch.func(以前称为“functorch”)是 类似 JAX 的 PyTorch 可组合函数转换。

注意

此库目前处于 测试阶段。这意味着这些功能通常有效(除非另有说明),并且我们(PyTorch 团队)致力于推进此库。但是,API 可能会根据用户反馈而更改,并且我们对 PyTorch 操作的覆盖范围不完整。

如果您对 API 或希望涵盖的用例有任何建议,请提交 GitHub 问题或与我们联系。我们很乐意了解您如何使用该库。

什么是可组合函数转换?

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

  • torch.func 具有自动微分转换(grad(f) 返回一个计算 f 梯度的函数)、向量化/批处理转换(vmap(f) 返回一个在输入批次上计算 f 的函数)等等。

  • 这些函数转换可以任意相互组合。例如,组合 vmap(grad(f)) 计算一个称为每样本梯度的量,而现有的 PyTorch 无法有效地计算。

为什么要使用可组合函数转换?

今天在 PyTorch 中有一些难以处理的用例

  • 计算每样本梯度(或其他每样本量)

  • 在一台机器上运行模型集成

  • 有效地将任务批处理到 MAML 的内循环中

  • 有效地计算雅可比矩阵和海森矩阵

  • 有效地计算批处理雅可比矩阵和海森矩阵

组合 vmap()grad()vjp() 转换使我们能够表达上述内容,而无需为每个内容设计单独的子系统。这种可组合函数转换的想法来自 JAX 框架

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发人员的深度教程

查看教程

资源

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

查看资源