元设备¶
“元”设备是一个抽象设备,表示仅记录元数据但不包含实际数据的张量。元张量有两个主要用例
模型可以在元设备上加载,允许您加载模型的表示,而无需实际将实际参数加载到内存中。如果您需要在加载实际数据之前对模型进行转换,这将非常有用。
大多数运算可以在元张量上执行,生成新的元张量,描述如果您在真实张量上执行该运算,结果将是什么。您可以使用它来执行抽象分析,而无需在计算或空间上花费时间来表示实际张量。由于元张量没有实际数据,因此您无法执行依赖数据的运算,例如
torch.nonzero()
或item()
。在某些情况下,并非所有设备类型(例如,CPU 和 CUDA)都对某一运算具有完全相同的输出元数据;在这种情况下,我们通常更倾向于忠实地表示 CUDA 行为。
警告
尽管原则上元张量计算应该始终比等效的 CPU/CUDA 计算更快,但许多元张量实现是在 Python 中实现的,尚未移植到 C++ 以提高速度,因此您可能会发现,使用小型 CPU 张量会获得更低的绝对框架延迟。
使用元张量的习惯用法¶
可以使用 torch.load()
将对象加载到元设备上,方法是指定 map_location='meta'
>>> torch.save(torch.randn(2), 'foo.pt')
>>> torch.load('foo.pt', map_location='meta')
tensor(..., device='meta', size=(2,))
如果您有一些任意代码,它执行一些张量构造,但没有明确指定设备,您可以使用 torch.device()
上下文管理器来覆盖它,以便改为在元设备上进行构造
>>> 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)
您无法直接将元张量转换为 CPU/CUDA 张量,因为元张量不存储任何数据,我们不知道您新张量的正确数据值是什么
>>> 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
包含用于获取任意张量并构造具有高保真度的等效元张量的未记录实用程序。这些 API 属于实验性质,可能会在任何时候以破坏向后兼容的方式更改。