快捷方式

Meta 设备

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

  • 模型可以加载到 meta 设备上,这允许您加载模型的表示而无需将实际参数加载到内存中。如果您需要在加载实际数据之前对模型进行转换,这将很有帮助。

  • 大多数操作可以在 meta 张量上执行,生成新的 meta 张量,这些张量描述了如果您在真实张量上执行该操作会得到的结果。您可以使用它来执行抽象分析,而无需花费计算时间或空间来表示实际张量。由于 meta 张量没有真实数据,您无法执行依赖于数据的操作,例如 torch.nonzero()item()。在某些情况下,并非所有设备类型(例如 CPU 和 CUDA)对于某个操作都具有完全相同的输出元数据;在这种情况下,我们通常倾向于忠实地表示 CUDA 的行为。

警告

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

使用 meta 张量的惯用法

通过指定 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,))

如果您有一些任意代码在未显式指定设备的情况下执行张量构建,您可以通过使用 torch.device() 上下文管理器来覆盖它,以便在 meta 设备上构建

>>> with torch.device('meta'):
...     print(torch.randn(30, 30))
...
tensor(..., device='meta', size=(30, 30))

这对于 NN 模块构建特别有用,在 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 包含未记录的实用程序,用于获取任意 Tensor 并构建具有高保真度的等效 meta Tensor。这些 API 是实验性的,随时可能以破坏 BC 的方式更改。

文档

查阅 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源