torch.func.vmap¶
- torch.func.vmap(func, in_dims=0, out_dims=0, randomness='error', *, chunk_size=None)[來源]¶
vmap 是向量化對映；
vmap(func)返回一個新的函式，該函式將func對映到輸入張量的某個維度上。在語義上，vmap 將對映操作推送到由func呼叫的 PyTorch 操作中，從而有效地向量化這些操作。vmap 對於處理批次維度非常有用：可以編寫一個在單個樣本上執行的函式
func，然後使用vmap(func)將其提升為可以接受批次樣本的函式。vmap 還可以與 autograd 結合使用，計算批次梯度。注意
torch.vmap()是torch.func.vmap()的別名，為了方便起見。你可以使用其中任何一個。- 引數
func (function) – 一個接受一個或多個引數的 Python 函式。必須返回一個或多個 Tensor。
in_dims (int 或 巢狀結構) – 指定輸入張量的哪個維度應被對映。
in_dims的結構應與輸入張量類似。如果特定輸入的in_dim為 None，則表明沒有對映維度。預設值：0。out_dims (int 或 Tuple[int]) – 指定對映維度應出現在輸出張量的哪個位置。如果
out_dims是一個 Tuple，則應包含與每個輸出對應的元素。預設值：0。randomness (str) – 指定此 vmap 中的隨機性在不同批次之間是相同還是不同。如果是 ‘different’，則每個批次的隨機性將不同。如果是 ‘same’，則隨機性在不同批次之間將相同。如果是 ‘error’，則呼叫任何隨機函式將報錯。預設值：‘error’。警告：此標誌僅適用於 PyTorch 的隨機操作，不適用於 Python 的 random 模組或 numpy 的隨機性。
chunk_size (None 或 int) – 如果為 None（預設值），則對輸入應用單個 vmap。如果不為 None，則每次計算
chunk_size個樣本的 vmap。注意，chunk_size=1等同於透過 for 迴圈計算 vmap。如果你在計算 vmap 時遇到記憶體問題，請嘗試使用非 None 的 chunk_size。
- 返回值
返回一個新的“批次化”函式。它接受與
func相同的輸入，不同之處在於每個輸入在由in_dims指定的索引處多了一個維度。它返回與func相同的輸出，不同之處在於每個輸出在由out_dims指定的索引處多了一個維度。- 返回型別
使用
vmap()的一個例子是計算批次點積。PyTorch 沒有提供批次化的torch.dotAPI；無需在文件中徒勞地翻找，可以使用vmap()來構造一個新的函式。>>> torch.dot # [D], [D] -> [] >>> batched_dot = torch.func.vmap(torch.dot) # [N, D], [N, D] -> [N] >>> x, y = torch.randn(2, 5), torch.randn(2, 5) >>> batched_dot(x, y)
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 >>> return feature_vec.dot(weights).relu() >>> >>> examples = torch.randn(batch_size, feature_size) >>> result = torch.vmap(model)(examples)
vmap()還可以幫助向量化之前難以或不可能批次計算的操作。一個例子是高階梯度計算。PyTorch 的 autograd 引擎計算 vjp（向量-雅可比積）。計算某個函式 f: R^N -> R^N 的完整雅可比矩陣通常需要呼叫 N 次autograd.grad，每次計算一行雅可比矩陣。使用vmap()，我們可以向量化整個計算過程，在一次呼叫autograd.grad中計算整個雅可比矩陣。>>> # Setup >>> N = 5 >>> f = lambda x: x ** 2 >>> x = torch.randn(N, requires_grad=True) >>> y = f(x) >>> I_N = torch.eye(N) >>> >>> # Sequential approach >>> jacobian_rows = [torch.autograd.grad(y, x, v, retain_graph=True)[0] >>> for v in I_N.unbind()] >>> jacobian = torch.stack(jacobian_rows) >>> >>> # vectorized gradient computation >>> def get_vjp(v): >>> return torch.autograd.grad(y, x, v) >>> jacobian = torch.vmap(get_vjp)(I_N)
vmap()還可以巢狀使用，生成具有多個批次維度的輸出>>> torch.dot # [D], [D] -> [] >>> batched_dot = torch.vmap(torch.vmap(torch.dot)) # [N1, N0, D], [N1, N0, D] -> [N1, N0] >>> x, y = torch.randn(2, 3, 5), torch.randn(2, 3, 5) >>> batched_dot(x, y) # tensor of size [2, 3]
如果輸入不是沿第一個維度進行批次化，
in_dims指定了每個輸入進行批次化的維度，例如>>> torch.dot # [N], [N] -> [] >>> batched_dot = torch.vmap(torch.dot, in_dims=1) # [N, D], [N, D] -> [D] >>> x, y = torch.randn(2, 5), torch.randn(2, 5) >>> batched_dot(x, y) # output is [5] instead of [2] if batched along the 0th dimension
如果存在多個輸入，且每個輸入沿不同維度進行批次化，則
in_dims必須是一個元組，其中包含每個輸入的批次維度，例如>>> torch.dot # [D], [D] -> [] >>> batched_dot = torch.vmap(torch.dot, in_dims=(0, None)) # [N, D], [D] -> [N] >>> x, y = torch.randn(2, 5), torch.randn(5) >>> batched_dot(x, y) # second arg doesn't have a batch dim because in_dim[1] was None
如果輸入是 Python 結構，
in_dims必須是包含與輸入形狀匹配的結構的元組>>> f = lambda dict: torch.dot(dict['x'], dict['y']) >>> x, y = torch.randn(2, 5), torch.randn(5) >>> input = {'x': x, 'y': y} >>> batched_dot = torch.vmap(f, in_dims=({'x': 0, 'y': None},)) >>> batched_dot(input)
預設情況下，輸出是沿第一個維度進行批次化。但是，可以透過使用
out_dims將其沿任何維度進行批次化>>> f = lambda x: x ** 2 >>> x = torch.randn(2, 5) >>> batched_pow = torch.vmap(f, out_dims=1) >>> batched_pow(x) # [5, 2]
對於任何使用 kwargs 的函式，返回的函式不會對 kwargs 進行批次化，但會接受 kwargs
>>> x = torch.randn([2, 5]) >>> def fn(x, scale=4.): >>> return x * scale >>> >>> batched_pow = torch.vmap(fn) >>> assert torch.allclose(batched_pow(x), x * 4) >>> batched_pow(x, scale=x) # scale is not batched, output has shape [2, 2, 5]
注意
vmap 不提供通用的自動批處理，也不能直接處理變長序列。
