torch.func 旋風之旅

什麼是 torch.func？

torch.func，先前稱為 functorch，是一個用於 PyTorch 中類似 JAX 的可組合函數轉換的程式庫。

• 「函數轉換」是一個高階函數，它接受一個數值函數，並傳回一個計算不同數量的新函數。

• torch.func 具有自動微分轉換（grad(f) 傳回一個計算 f 梯度函數），一個向量化/批次處理轉換（vmap(f) 傳回一個計算 f 在輸入批次上的函數），以及其他。

• 這些函數轉換可以任意互相組合。例如，組合 vmap(grad(f)) 會計算一個稱為每樣本梯度 (per-sample-gradients) 的數量，這是目前標準 PyTorch 無法有效計算的。

為什麼要使用可組合的函數轉換？

目前在 PyTorch 中，有一些使用案例很難實現：- 計算每個樣本的梯度（或其他每個樣本的量）

• 在單一機器上執行模型集成

• 有效地將 MAML 內部迴圈中的任務批次處理在一起

• 有效率地計算 Jacobian 和 Hessian 矩陣

• 有效率地計算批次化的 Jacobian 和 Hessian 矩陣

組合 vmap()、grad()、vjp() 和 jvp() 轉換，讓我們能夠表達以上功能，而無需為每個功能設計一個單獨的子系統。

什麼是轉換？

grad() (梯度計算)

grad(func) 是我們的梯度計算轉換。它會回傳一個新的函式，該函式會計算 func 的梯度。它假設 func 回傳一個單一元素的 Tensor，並且預設情況下，它會計算 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 中的所有 Tensor 運算添加一個維度。vmap(func) 回傳一個新的函式，該函式會將 func 映射到 inputs 中每個 Tensor 的某個維度（預設值：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):
    # 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() (向量-Jacobian 乘積)

vjp() 轉換將 func 應用於 inputs，並回傳一個新的函式，該函式會計算給定一些 cotangents Tensors 的向量-Jacobian 乘積 (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() (Jacobian-向量乘積)

jvp() 轉換計算 Jacobian-向量乘積，也稱為“正向模式 AD”。與大多數其他轉換不同，它不是高階函式，但它會回傳 func(inputs) 的輸出以及 jvps。

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

jacrev() 轉換回傳一個新的函式，該函式接收 x，並使用反向模式 AD 回傳函式相對於 x 的 Jacobian 矩陣。

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() 組合使用，以產生批次化的 Jacobian 矩陣

x = torch.randn(64, 5)
jacobian = vmap(jacrev(torch.sin))(x)
assert jacobian.shape == (64, 5, 5)

jacfwd() 是 jacrev 的直接替代品，它使用正向模式 AD 計算 Jacobian 矩陣

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() 組合使用可以產生 Hessian 矩陣

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)