torch.vmap

torch.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 (函式) – 一個接受一個或多個引數的 Python 函式。必須返回一個或多個張量。

• 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 指定的索引處有一個額外的維度。

返回型別

Callable

使用 vmap() 的一個示例是計算批處理點積。PyTorch 沒有提供批處理的 torch.dot API；與其徒勞地翻找文件，不如使用 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 引擎計算 vjps（向量-雅可比積）。計算某個函式 f: R^N -> R^N 的完整雅可比矩陣通常需要對 autograd.grad 進行 N 次呼叫，每行雅可比矩陣呼叫一次。使用 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 不提供通用的自動批處理功能，也不直接處理變長序列。