元设备¶
“元”设备是一个抽象设备,表示仅记录元数据但没有实际数据的张量。元张量有两个主要用例
模型可以加载到元设备上,允许您加载模型的表示,而无需实际将参数加载到内存中。如果您需要在加载实际数据之前对模型进行转换,这将非常有用。
大多数操作都可以在元张量上执行,生成新的元张量,这些张量描述了如果您在真实张量上执行操作会得到什么结果。您可以使用此功能执行抽象分析,而无需花费时间在计算或空间上表示实际张量。由于元张量没有真实数据,因此您无法执行
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 处于实验阶段,可能在任何时候以破坏性兼容的方式发生更改。