Table of Contents
Shortcuts

torch.func 快速入门

什么是 torch.func？

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

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

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

  • 这些函数变换可以任意地相互组合。例如，组合 vmap(grad(f)) 计算了一个称为每样本梯度的量，这是 PyTorch 目前无法高效计算的。

为什么使用可组合的函数变换？

在 PyTorch 中，有一些用例目前做起来比较棘手： - 计算每个样本的梯度（或其他每个样本的量）

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

  • 在内循环中高效地将任务批处理在一起

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

  • 高效计算批量雅可比矩阵和海森矩阵

组合 vmap(), grad(), vjp(), 和 jvp() 变换 使我们能够表达上述内容，而无需为每个内容设计一个单独的子系统。

什么是变换？

grad()（梯度计算）

grad(func) 是我们的梯度计算变换。它返回一个新函数，用于计算 func 的梯度。它假设 func 返回一个单元素张量，并且默认情况下，它计算 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())

# 二阶梯度
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 torch.func import vmap
batch_size, feature_size = 3, 5
weights = torch.randn(feature_size, requires_grad=True)

def model(feature_vec):
    # 非常简单的线性模型，带有激活函数
    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):
    # 非常简单的线性模型，带有激活函数
    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()（向量-雅可比积）

The vjp() transform applies func to inputs and returns a new function that computes the vector-Jacobian product (vjp) given some cotangents Tensors.

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()（雅可比向量积）

The jvp() 变换计算雅可比向量积，也被称为“前向模式自动微分”。与大多数其他变换不同，它不是一个高阶函数，但它返回 func(inputs) 的输出以及雅可比向量积。

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()

变换返回一个新函数，该函数接受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 的直接替代品，使用前向模式自动微分计算雅可比矩阵：

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() 可以生成海森矩阵：

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)
优云智算