快捷方式

Meta 裝置

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

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

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

警告

儘管原則上 meta 張量計算應始終比等效的 CPU/CUDA 計算更快,但許多 meta 張量實現在 Python 中實現,尚未移植到 C++ 以提高速度,因此您可能會發現在使用小型 CPU 張量時獲得較低的絕對框架延遲。

使用 meta 張量的慣用法

透過指定 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,))

如果您有一些任意程式碼在未顯式指定裝置的情況下執行張量構建,您可以透過使用 torch.device() 上下文管理器來覆蓋它,以便在 meta 裝置上構建

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

這對於 NN 模組構建特別有用,在 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 張量直接轉換為 CPU/CUDA 張量,因為 meta 張量不儲存資料,我們不知道新張量的正確資料值是什麼

>>> 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 是實驗性的,隨時可能以破壞 BC 的方式更改。

文件

查閱 PyTorch 的全面開發者文件

檢視文件

教程

獲取面向初學者和高階開發者的深度教程

檢視教程

資源

查詢開發資源並獲得問題解答

檢視資源