Meta device#

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

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

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

警告

虽然理论上 meta tensor 的计算应该总是比等效的 CPU/CUDA 计算更快，但许多 meta tensor 实现是用 Python 编写的，尚未移植到 C++ 以提高速度，因此你可能会发现对于小型 CPU tensor，绝对框架延迟更低。

使用 meta tensor 的惯用法#

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

>>> torch.save(torch.randn(2), 'foo.pt')
>>> torch.load('foo.pt', map_location='meta')
tensor(..., device='meta', size=(2,))

如果你有一些不明确指定设备的任意代码来执行 tensor 构建，你可以使用 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 tensor 转换为 CPU/CUDA tensor，因为 meta tensor 不存储任何数据，我们不知道新 tensor 的正确数据值是多少。

>>> 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 并以高保真度构建等效的 Meta Tensor。这些 API 是实验性的，可能随时以向后不兼容的方式进行更改。

Meta device#

使用 meta tensor 的惯用法#

文档

教程

资源