Meta 设备#

创建日期：2025 年 6 月 17 日 | 最后更新日期：2025 年 6 月 17 日

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

模型可以加载到 meta 设备上，这样您就可以加载模型的表示形式，而无需将实际参数加载到内存中。如果您需要在加载实际数据之前对模型进行转换，这会很有帮助。
大多数操作都可以在 meta 张量上执行，生成新的 meta 张量，描述如果您在实际张量上执行操作，结果会是什么。您可以使用它进行抽象分析，而无需花费计算时间或空间来表示实际张量。由于 meta 张量没有实际数据，您无法执行依赖数据的操作，例如 torch.nonzero() 或 item()。在某些情况下，并非所有设备类型（例如 CPU 和 CUDA）对某个操作都具有完全相同的输出元数据；在这种情况下，我们通常更倾向于忠实地表示 CUDA 行为。

警告

尽管原则上 meta 张量计算应该始终比等效的 CPU/CUDA 计算更快，但许多 meta 张量实现在 Python 中实现，尚未移植到 C++ 以提高速度，因此您可能会发现使用小型 CPU 张量时，绝对框架延迟会更低。

使用 meta 张量的惯用语#

可以通过指定 map_location='meta' 将对象加载到 meta 设备上，使用 torch.load()。

>>> 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 模块的构造特别有帮助，因为您通常无法明确传入设备进行初始化。

>>> 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 包含未记录的实用程序，用于获取任意张量并以高保真度构造等效的 meta 张量。这些 API 处于实验阶段，可能会随时以不向后兼容的方式更改。

Meta 设备#

使用 meta 张量的惯用语#

文档

教程

资源