模型集成¶
此範例說明如何使用 vmap 對模型集成進行向量化。
什麼是模型集成？¶
模型集成是將多個模型的預測結果組合在一起。傳統上，這是透過分別在某些輸入上運行每個模型，然後組合預測結果來完成的。但是，如果您要運行具有相同架構的模型，則可以使用 vmap 將它們組合在一起。 vmap 是一種函數轉換，它可以跨輸入張量的維度映射函數。它的用例之一是消除 for 迴圈，並透過向量化來加速它們。
讓我們演示如何使用簡單 MLP 的集成來做到這一點。
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
torch.manual_seed(0);
# Here's a simple MLP
class SimpleMLP(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(784, 128)
self.fc2 = nn.Linear(128, 128)
self.fc3 = nn.Linear(128, 10)
def forward(self, x):
x = x.flatten(1)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
x = F.relu(x)
x = self.fc3(x)
return x
讓我們生成一批虛擬數據，並假設我們正在使用 MNIST 數據集。因此，虛擬圖像是 28 x 28，我們有一個大小為 64 的 minibatch。此外，假設我們要組合來自 10 個不同模型的預測結果。
device = 'cuda'
num_models = 10
data = torch.randn(100, 64, 1, 28, 28, device=device)
targets = torch.randint(10, (6400,), device=device)
models = [SimpleMLP().to(device) for _ in range(num_models)]
我們有幾個選項來生成預測結果。也許我們想給每個模型一個不同的隨機 minibatch 數據。或者，也許我們想透過每個模型運行相同的 minibatch 數據（例如，如果我們正在測試不同模型初始化的效果）。
選項 1：每個模型使用不同的 minibatch
minibatches = data[:num_models]
predictions_diff_minibatch_loop = [model(minibatch) for model, minibatch in zip(models, minibatches)]
選項 2：相同的 minibatch
minibatch = data[0]
predictions2 = [model(minibatch) for model in models]
使用 vmap 對集成進行向量化¶
讓我們使用 vmap 來加速 for 迴圈。我們必須首先準備好模型以供 vmap 使用。
首先，讓我們透過堆疊每個參數來組合模型的狀態。例如，model[i].fc1.weight 的形狀為 [784, 128]；我們要將 10 個模型的 .fc1.weight 堆疊起來，產生一個形狀為 [10, 784, 128] 的大權重。
functorch 提供了「combine_state_for_ensemble」便捷函數來做到這一點。它返回模型的無狀態版本 (fmodel) 以及堆疊的參數和緩衝區。
from functorch import combine_state_for_ensemble
fmodel, params, buffers = combine_state_for_ensemble(models)
[p.requires_grad_() for p in params];
選項 1：使用每個模型的不同 minibatch 獲取預測結果。
預設情況下，vmap 會將函數映射到傳遞給函數的所有輸入的第一個維度。使用 combine_state_for_ensemble 後，每個參數和緩衝區在前面都有一個大小為「num_models」的額外維度，而 minibatch 有一個大小為「num_models」的維度。
print([p.size(0) for p in params]) # show the leading 'num_models' dimension
assert minibatches.shape == (num_models, 64, 1, 28, 28) # verify minibatch has leading dimension of size 'num_models'
[10, 10, 10, 10, 10, 10]
from functorch import vmap
predictions1_vmap = vmap(fmodel)(params, buffers, minibatches)
# verify the vmap predictions match the
assert torch.allclose(predictions1_vmap, torch.stack(predictions_diff_minibatch_loop), atol=1e-3, rtol=1e-5)
選項 2：使用相同的 minibatch 數據獲取預測結果。
vmap 有一個 in_dims 參數，用於指定要映射的維度。透過使用 None，我們告訴 vmap 我們希望將相同的 minibatch 應用於所有 10 個模型。
predictions2_vmap = vmap(fmodel, in_dims=(0, 0, None))(params, buffers, minibatch)
assert torch.allclose(predictions2_vmap, torch.stack(predictions2), atol=1e-3, rtol=1e-5)
簡要說明：vmap 可以轉換的函數類型有限制。最適合轉換的函數是純函數：輸出僅由輸入確定的函數，沒有副作用（例如突變）。vmap 無法處理任意 Python 數據結構的突變，但它能夠處理許多就地 PyTorch 操作。
效能¶
對效能數字感到好奇嗎？以下是 Google Colab 上的數字。
from torch.utils.benchmark import Timer
without_vmap = Timer(
stmt="[model(minibatch) for model, minibatch in zip(models, minibatches)]",
globals=globals())
with_vmap = Timer(
stmt="vmap(fmodel)(params, buffers, minibatches)",
globals=globals())
print(f'Predictions without vmap {without_vmap.timeit(100)}')
print(f'Predictions with vmap {with_vmap.timeit(100)}')
Predictions without vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7fe22c58b3d0>
[model(minibatch) for model, minibatch in zip(models, minibatches)]
3.25 ms
1 measurement, 100 runs , 1 thread
Predictions with vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7fe22c50c450>
vmap(fmodel)(params, buffers, minibatches)
879.28 us
1 measurement, 100 runs , 1 thread
使用 vmap 後速度大幅提升！
一般來說，使用 vmap 進行向量化應該比在 for 迴圈中運行函數更快，並且與手動批次處理相比具有競爭力。但也有一些例外情況，例如，如果我們沒有為特定操作實作 vmap 規則，或者如果基礎內核沒有針對舊硬體（GPU）進行優化。如果您遇到任何這些情況，請透過在我們的 GitHub 上提交 issue 告知我們！
