元设备 ¶

“元”设备是一种抽象设备，表示一个仅记录元数据而不记录实际数据的张量。元张量有两个主要用途：

• 模型可以加载到元设备上，允许您在不实际将参数加载到内存中的情况下加载模型的表示。如果您需要在加载实际数据之前对模型进行转换，这可能会很有帮助。

• 大多数操作可以在元张量上执行，生成新的元张量，描述如果对真实张量执行该操作的结果。您可以使用此功能进行抽象分析，而无需花费时间在计算或空间上表示实际张量。因为元张量没有真实数据，所以不能执行依赖于数据的操作，如 torch.nonzero() 或 item() 。在某些情况下，并非所有设备类型（例如 CPU 和 CUDA）的操作输出元数据都完全相同；在这种情况下，我们通常更喜欢忠实反映 CUDA 的行为。

警告

虽然在原则上元张量计算应该始终比等效的 CPU/CUDA 计算更快，但许多元张量实现都是用 Python 编写的，尚未移植到 C++以提高速度，因此您可能会发现，对于小的 CPU 张量，您得到的绝对框架延迟更低。

元张量操作习语

可以通过指定 map_location='meta' 将对象加载到元设备 torch.load() ：

>>> 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))

这对于神经网络模块的构造尤其有帮助，因为在初始化时，你通常无法明确传递设备：

>>> 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 包含用于将任意 Tensor 转换为等效元 Tensor 的未记录实用工具，这些 API 是实验性的，并且可能在任何时间以 BC 破坏的方式更改。