目录

Meta设备

“meta”设备是一种抽象设备,表示一个仅记录元数据而不包含实际数据的张量。Meta 张量主要有两种应用场景:

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

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

警告

尽管原则上元张量计算应该总是比等效的CPU/CUDA计算更快,但许多元张量实现是在Python中进行的,并且尚未移植到C++以提高速度,因此你会发现对于小规模的CPU张量,你可能会得到更低的绝对框架延迟。

处理元张量的惯用法

一个对象可以通过指定 torch.load() map_location='meta'加载到元设备上:

>>> 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 包含未记录的工具,用于将任意张量转换为具有高保真的等效元张量。这些API是实验性的,并且可能在任何时候以破坏向后兼容性的方式更改。

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

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

查看教程

资源

查找开发资源并解答您的问题

查看资源