評價此頁

Meta device#

建立日期:2025 年 6 月 17 日 | 最後更新日期:2025 年 6 月 17 日

“meta” 裝置是一個抽象裝置,用於表示一個只記錄元資料而不記錄實際資料的 tensor。Meta tensor 有兩個主要用例:

  • 可以在 meta 裝置上載入模型,這樣你就可以載入模型的表示,而無需將實際引數載入到記憶體中。如果你需要在載入實際資料之前對模型進行轉換,這會很有幫助。

  • 大多數操作都可以在 meta tensor 上執行,生成新的 meta tensor,描述如果你在真實 tensor 上執行該操作,結果會是什麼樣子。你可以使用它來進行抽象分析,而無需花費時間和空間來表示實際的 tensor。由於 meta tensor 沒有真實資料,因此無法執行依賴於資料的操作,例如 torch.nonzero()item()。在某些情況下,並非所有裝置型別(例如 CPU 和 CUDA)對某個操作都具有完全相同的輸出元資料;在這種情況下,我們通常傾向於忠實地表示 CUDA 的行為。

警告

雖然理論上 meta tensor 的計算應該總是比等效的 CPU/CUDA 計算更快,但許多 meta tensor 實現是用 Python 編寫的,尚未移植到 C++ 以提高速度,因此你可能會發現對於小型 CPU tensor,絕對框架延遲更低。

使用 meta tensor 的慣用法#

可以透過指定 map_location='meta' 來使用 torch.load() 將物件載入到 meta 裝置。

>>> torch.save(torch.randn(2), 'foo.pt')
>>> torch.load('foo.pt', map_location='meta')
tensor(..., device='meta', size=(2,))

如果你有一些不明確指定裝置的任意程式碼來執行 tensor 構建,你可以使用 torch.device() 上下文管理器來覆蓋它,使其改為在 meta 裝置上構建。

>>> with torch.device('meta'):
...     print(torch.randn(30, 30))
...
tensor(..., device='meta', size=(30, 30))

這對於 NN 模組的構建尤其有用,因為你通常無法在初始化時明確傳遞裝置。

>>> from torch.nn.modules import Linear
>>> with torch.device('meta'):
...     print(Linear(20, 30))
...
Linear(in_features=20, out_features=30, bias=True)

你無法直接將 meta tensor 轉換為 CPU/CUDA tensor,因為 meta tensor 不儲存任何資料,我們不知道新 tensor 的正確資料值是多少。

>>> torch.ones(5, device='meta').to("cpu")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
NotImplementedError: Cannot copy out of meta tensor; no data!

使用像 torch.empty_like() 這樣的工廠函式來顯式指定如何填充缺失的資料。

NN 模組有一個方便的方法 torch.nn.Module.to_empty(),它允許你將模組移動到另一個裝置,同時將所有引數保持未初始化狀態。你需要手動顯式重新初始化引數。

>>> from torch.nn.modules import Linear
>>> with torch.device('meta'):
...     m = Linear(20, 30)
>>> m.to_empty(device="cpu")
Linear(in_features=20, out_features=30, bias=True)

torch._subclasses.meta_utils 包含一些未公開的實用程式,用於接收任意 Tensor 並以高保真度構建等效的 Meta Tensor。這些 API 是實驗性的,可能隨時以向後不相容的方式進行更改。