torch.func 快速入門

建立日期：2025 年 6 月 12 日 | 最後更新日期：2025 年 6 月 12 日

什麼是 torch.func？

torch.func，以前稱為 functorch，是一個用於 PyTorch 中實現 JAX 式可組合函式變換的庫。

• “函式變換”是一種高階函式，它接受一個數值函式並返回一個計算不同量的新函式。

• torch.func 包含自動微分變換（grad(f) 返回一個計算 f 的梯度的函式）、向量化/批處理變換（vmap(f) 返回一個在輸入批次上計算 f 的函式）等。

• 這些函式變換可以任意組合。例如，組合 vmap(grad(f)) 可以計算一個稱為“每樣本梯度”的量，這是標準 PyTorch 目前無法高效計算的。

為什麼需要可組合函式變換？

目前在 PyTorch 中有一些難以實現的使用案例

• 計算逐樣本梯度（或其他逐樣本量）

• 在單機上執行模型整合

• 高效地批處理 MAML 內部迴圈中的任務

• 高效地計算雅可比矩陣和 Hessian 矩陣

• 高效地計算批處理雅可比矩陣和 Hessian 矩陣

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

# Second-order gradients
neg_sin_x = grad(grad(lambda x: torch.sin(x)))(x)
assert torch.allclose(neg_sin_x, -x.sin())

vmap()（自動向量化）

注意：vmap() 對其可使用的程式碼施加了限制。有關更多詳細資訊，請參閱 UX 限制。

vmap(func)(*inputs) 是一個變換，它為 func 中的所有張量運算添加了一個維度。vmap(func) 返回一個新函式，該函式將在輸入中每個張量的某個維度（預設為 0）上對映 func。

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):
    # 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 torch.func 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()（向量-雅可比積）

給定某些 cotangents 張量，vjp() 變換將 func 應用於 inputs，並返回一個計算向量-雅可比積 (vjp) 的新函式。

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

jvp() 變換計算雅可比-向量積，也稱為“前向模式 AD”。與大多數其他變換不同，它不是高階函式，但它會返回 func(inputs) 的輸出以及 jvp。

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

使用反向模式 AD，jacrev() 變換返回一個新函式，該函式接受 x 並返回函式相對於 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)

使用前向模式 AD 計算雅可比矩陣時，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)