元设备

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

• 模型可以在元设备上加载，允许您加载模型的表示形式，而无需实际将模型参数加载到内存中。如果您需要在加载实际数据之前对模型进行转换，这将非常有用。

• 大多数操作可以在元张量上执行，生成新的元张量，描述如果在真实张量上执行操作会得到什么结果。您可以使用它来执行抽象分析，而无需花费时间在计算或空间上表示实际张量。由于元张量没有真实数据，因此无法执行依赖数据的操作，例如 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))

这在 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 处于实验阶段，可能会在任何时候以破坏向后兼容的方式更改。