Meta 设备¶
“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 模块构建尤其有帮助,在 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 是实验性的,可能会随时以不兼容 BC 的方式进行更改。