torch.vmap¶
- torch.vmap(func, in_dims=0, out_dims=0, randomness='error', *, chunk_size=None)[source]¶
vmap 是向量化的 map；
vmap(func)會回傳一個新的函式，將func對輸入的某些維度進行映射。從語義上來說，vmap 會將 map 推入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，則表示沒有 map 維度。預設值：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，則一次計算 vmap
chunk_size個樣本。 請注意，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 引擎計算 vjps（向量-雅可比矩陣乘積）。 對於某些函數 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 不提供通用的自動批次處理或直接處理可變長度的序列。
