Shortcuts

元设备

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

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

  • 大多数操作都可以在元张量上执行,生成新的元张量,这些元张量描述了如果在实际张量上执行操作会得到的结果。你可以使用这个功能来进行抽象分析,而不需要花费时间在计算或空间上来表示实际的张量。因为元张量没有实际数据,所以你不能执行依赖于数据的操作,例如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')
张量(..., 设备='meta', 大小=(2,))

如果你有一些任意代码在执行张量构造时没有明确指定设备,你可以通过使用torch.device()上下文管理器来覆盖它,从而在meta设备上进行构造:

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

这在NN模块构建中特别有用,因为在这种情况下,您通常无法显式地传入用于初始化的设备:

>>> from torch.nn.modules import Linear
>>> with torch.device('meta'):
...     print(Linear(20, 30))
...
线性(in_features=20, out_features=30, bias=True)

您不能直接将元张量转换为CPU/CUDA张量,因为元张量不存储数据,我们不知道新张量的正确数据值是什么:

>>> torch.ones(5, device='meta').to("cpu")
Traceback (most recent call last):
  File "", line 1, in 
NotImplementedError: 无法从元张量复制数据;没有数据!

使用像 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")
线性(in_features=20, out_features=30, bias=True)

torch._subclasses.meta_utils 包含用于获取任意张量并构建具有高保真度的等效元张量的未记录实用程序。这些API是实验性的,可能会在任何时候以破坏向后兼容性的方式进行更改。

优云智算