旋風式導覽¶
什麼是 functorch？¶
functorch 是一個用於在 PyTorch 中進行類似 JAX 的可組合函數變換的函式庫。
「函數變換」是一種高階函數，它接受一個數值函數並返回一個計算不同數量的新函數。
functorch 具有自動微分變換（
grad(f)返回一個計算f的梯度的函數）、向量化/批次處理變換（vmap(f)返回一個計算f在批次輸入上的函數）等等。這些函數變換可以任意組合。例如，組合
vmap(grad(f))計算一個稱為每個樣本梯度的量，而目前的 PyTorch 無法有效地計算它。
此外，我們還在 functorch.compile 命名空間中提供了一個實驗性的編譯變換。我們的編譯變換稱為 AOT（預先）Autograd，它會返回一個 FX 圖（可選地包含反向傳播），您可以透過各種後端對其進行編譯。
為什麼要使用可組合函數變換？¶
在 PyTorch 中，有許多用例很難處理
計算每個樣本的梯度（或其他每個樣本的量）
在單台機器上運行模型集成
在 MAML 的內部循環中有效地批次處理任務
有效地計算雅可比矩陣和黑塞矩陣
有效地計算批次的雅可比矩陣和黑塞矩陣
組合 vmap、grad、vjp 和 jvp 變換，使我們能夠在不為每個子系統設計單獨的系統的情況下表達上述內容。
有哪些變換？¶
grad（梯度計算）¶
grad(func) 是我們的梯度計算變換。它返回一個計算 func 的梯度的新函數。它假設 func 返回一個單元素張量，並且默認情況下它計算 func 的輸出相對於第一個輸入的梯度。
import torch
from functorch 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 對其可使用的程式碼施加了限制。有關更多詳細資訊，請閱讀其文件字串。
vmap(func)(*inputs) 是一個變換，它為 func 中的所有張量操作添加一個維度。vmap(func) 返回一個新函數，該函數將 func 映射到輸入中每個張量的某個維度（默認值：0）。
vmap 對於隱藏批次維度很有用：可以編寫一個在範例上運行的函數 func，然後使用 vmap(func) 將其提升為可以處理批次範例的函數，從而簡化建模體驗
import torch
from functorch 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 functorch 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（向量-雅可比矩陣積）¶
vjp 變換將 func 應用於 inputs，並返回一個新函數，該函數在給定一些 cotangents 張量的情況下計算向量-雅可比矩陣積 (vjp)。
from functorch 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 變換計算雅可比矩陣-向量積，也稱為「正向模式自動微分」。與大多數其他變換不同，它不是高階函數，但它會返回 func(inputs) 的輸出以及 jvp。
from functorch import jvp
x = torch.randn(5)
y = torch.randn(5)
f = lambda x, y: (x * y)
_, output = jvp(f, (x, y), (torch.ones(5), torch.ones(5)))
assert torch.allclose(output, x + y)
jacrev、jacfwd 和 hessian¶
jacrev 變換返回一個新函數，該函數接受 x 並使用反向模式自動微分返回函數相對於 x 的雅可比矩陣。
from functorch 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 functorch 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 functorch import hessian
def f(x):
return x.sin().sum()
x = torch.randn(5)
hess = hessian(f)(x)
