torch.Storage¶
在 PyTorch 中,一个常规的 tensor 是一个多维数组,由以下组件定义:
Storage:tensor 的实际数据,存储为一个连续的一维字节数组。
dtype
:tensor 中元素的 数据类型,例如 torch.float32 或 torch.int64。shape
:一个元组,指示 tensor 在每个维度上的大小。Stride:在每个维度上从一个元素移动到下一个元素所需的步长。
Offset:存储中 tensor 数据开始的起始点。对于新创建的 tensor,这通常为 0。
这些组件共同定义了 tensor 的结构和数据,其中 storage 存储实际数据,其余部分作为元数据。
无类型 Storage API¶
torch.UntypedStorage
是一个连续的一维元素数组。其长度等于 tensor 的字节数。storage 作为 tensor 的底层数据容器。通常,在 PyTorch 中使用常规构造函数(如 zeros()
、zeros_like()
或 new_zeros()
)创建的 tensor,其 tensor storage 与 tensor 本身之间存在一对一的对应关系。
然而,一个 storage 可以被多个 tensor 共享。例如,tensor 的任何视图(通过 view()
或某些(但非全部)索引类型(如整数和切片)获得)将指向与原始 tensor 相同的底层 storage。序列化和反序列化共享同一 storage 的 tensor 时,这种关系得以保留,并且这些 tensor 继续指向同一 storage。有趣的是,反序列化指向单个 storage 的多个 tensor 比反序列化多个独立的 tensor 更快。
可以通过 untyped_storage()
方法访问 tensor storage。这将返回类型为 torch.UntypedStorage
的对象。幸运的是,storage 有一个唯一的标识符,通过 torch.UntypedStorage.data_ptr()
方法访问。在常规设置中,具有相同数据 storage 的两个 tensor 将具有相同的 storage data_ptr
。然而,tensor 本身可以指向两个独立的 storage,一个用于其 数据
属性,另一个用于其 梯度
属性。每个都需要自己的 data_ptr()
。通常,不能保证 torch.Tensor.data_ptr()
和 torch.UntypedStorage.data_ptr()
匹配,也不应假定如此。
无类型 storage 与基于它们的 tensor 有些独立。实际上,这意味着具有不同 dtype 或 shape 的 tensor 可以指向同一个 storage。这也意味着 tensor storage 可以被更改,如下例所示:
>>> t = torch.ones(3)
>>> s0 = t.untyped_storage()
>>> s0
0
0
128
63
0
0
128
63
0
0
128
63
[torch.storage.UntypedStorage(device=cpu) of size 12]
>>> s1 = s0.clone()
>>> s1.fill_(0)
0
0
0
0
0
0
0
0
0
0
0
0
[torch.storage.UntypedStorage(device=cpu) of size 12]
>>> # Fill the tensor with a zeroed storage
>>> t.set_(s1, storage_offset=t.storage_offset(), stride=t.stride(), size=t.size())
tensor([0., 0., 0.])
警告
请注意,直接修改 tensor 的 storage(如本例所示)不是推荐的做法。这种低级操作仅出于教育目的进行演示,以展示 tensor 及其底层 storage 之间的关系。通常,使用标准的 torch.Tensor
方法(例如 clone()
和 fill_()
)来实现相同结果更有效、更安全。
除了 data_ptr
,无类型 storage 还有其他属性,例如 filename
(如果 storage 指向磁盘上的文件)、device
或 is_cuda
用于设备检查。storage 还可以使用诸如 copy_
、fill_
或 pin_memory
之类的方法进行原地或非原地操作。有关更多信息,请查看下面的 API 参考。请记住,修改 storage 是一种低级 API,伴随着风险!大多数这些 API 也存在于 tensor 级别:如果存在,应优先使用它们的 tensor 对应项。
特殊情况¶
我们提到,具有非 None 梯度
属性的 tensor 实际上包含两个数据块。在这种情况下,untyped_storage()
将返回 数据
属性的 storage,而梯度的 storage 可以通过 tensor.grad.untyped_storage()
获取。
>>> t = torch.zeros(3, requires_grad=True)
>>> t.sum().backward()
>>> assert list(t.untyped_storage()) == [0] * 12 # the storage of the tensor is just 0s
>>> assert list(t.grad.untyped_storage()) != [0] * 12 # the storage of the gradient isn't
- 也存在 tensor 没有典型 storage 或根本没有 storage 的特殊情况:
"meta"
设备上的 Tensor:"meta"
设备上的 tensor 用于 shape 推断,不包含实际数据。Fake Tensors:PyTorch 编译器使用的另一个内部工具是 FakeTensor,它基于类似的想法。
Tensor 子类或类似 tensor 的对象也可能表现出异常行为。一般来说,我们不期望很多使用场景需要操作到 Storage 级别!
- class torch.UntypedStorage(*args, **kwargs)[source][source]¶
-
- copy_()¶
- cuda(device=None, non_blocking=False)[source]¶
返回此对象在 CUDA 内存中的副本。
如果此对象已在 CUDA 内存中且位于正确的设备上,则不执行复制并返回原始对象。
- 参数
- 返回类型
Union[_StorageBase, TypedStorage]
- data_ptr()¶
- element_size()¶
- property filename: Optional[str]¶
返回与此 storage 关联的文件名。
如果 storage 在 CPU 上并通过
from_file()
设置shared
为True
创建,则文件名将是一个字符串。否则,此属性为None
。
- fill_()¶
- static from_buffer()¶
- static from_file(filename, shared=False, size=0) Storage ¶
创建一个由内存映射文件支持的 CPU storage。
如果
shared
为True
,则所有进程之间共享内存。所有更改都会写入文件。如果shared
为False
,则 storage 上的更改不会影响文件。size
是 storage 中的元素数量。如果shared
为False
,则文件必须包含至少size * sizeof(Type)
字节(Type
是 storage 的类型,对于UnTypedStorage
,文件必须包含至少size
字节)。如果shared
为True
,如果需要将创建文件。- 参数
filename (str) – 要映射的文件名
shared (bool) – 是否共享内存(是传递
MAP_SHARED
还是MAP_PRIVATE
给底层 mmap(2) 调用)size (int) – storage 中的元素数量
- hpu(device=None, non_blocking=False)[source]¶
返回此对象在 HPU 内存中的副本。
如果此对象已在 HPU 内存中且位于正确的设备上,则不执行复制并返回原始对象。
- 参数
- 返回类型
Union[_StorageBase, TypedStorage]
- property is_cuda¶
- property is_hpu¶
- is_pinned(device='cuda')[source]¶
确定 CPU storage 是否已固定在设备上。
- 参数
device (str 或 torch.device) – 要固定内存的设备(默认值:
'cuda'
)。不建议使用此参数,它可能会被弃用。- 返回
一个布尔变量。
- nbytes()¶
- new()¶
- pin_memory(device='cuda')[source]¶
将 CPU 存储复制到固定内存,如果尚未固定。
- 参数
device (str 或 torch.device) – 要固定内存的设备(默认值:
'cuda'
)。不建议使用此参数,它可能会被弃用。- 返回
一个固定内存的 CPU 存储。
- resizable()¶
- resize_()¶
将存储移动到共享内存。
对于已在共享内存中的存储以及 CUDA 存储(它们无需移动即可跨进程共享),这是一个空操作。共享内存中的存储无法调整大小。
请注意,为缓解 this 等问题,在同一对象上从多个线程调用此函数是线程安全的。但是,在未进行适当同步的情况下,调用 self 上的任何其他函数都是非线程安全的。请参阅 Multiprocessing best practices 以获取更多详细信息。
注意
当共享内存中存储的所有引用都被删除时,关联的共享内存对象也将被删除。PyTorch 有一个特殊的清理流程,以确保即使当前进程意外退出,此操作也会发生。
值得注意的是
share_memory_()
与from_file()
并设置shared = True
之间的区别。share_memory_
使用 shm_open(3) 创建一个 POSIX 共享内存对象,而from_file()
使用 open(2) 打开用户传入的文件名。两者都使用带有
MAP_SHARED
的 mmap(2) 调用,将文件/对象映射到当前的虚拟地址空间。share_memory_
会在映射对象后调用shm_unlink(3)
,以确保当没有进程打开该对象时,共享内存对象会被释放。torch.from_file(shared=True)
不会取消链接该文件。此文件是持久的,将保留直到用户删除它。
- 返回
self
- type(dtype=None, non_blocking=False)[source]¶
- 返回类型
Union[_StorageBase, TypedStorage]
传统类型存储¶
警告
从历史角度来看,PyTorch 以前使用类型化存储类,这些类现已弃用,应避免使用。以下详细介绍了此 API,以防您遇到它,但强烈不建议使用它。将来,除了 torch.UntypedStorage
之外的所有存储类都将被移除,并且所有情况下都将使用 torch.UntypedStorage
。
torch.Storage
是对应于默认数据类型 (torch.get_default_dtype()
) 的存储类的别名。例如,如果默认数据类型是 torch.float
,则 torch.Storage
解析为 torch.FloatStorage
。
torch.<type>Storage
和 torch.cuda.<type>Storage
类,例如 torch.FloatStorage
、torch.IntStorage
等,实际上从未实例化。调用它们的构造函数会创建一个具有适当 torch.dtype
和 torch.device
的 torch.TypedStorage
。torch.<type>Storage
类拥有 torch.TypedStorage
所拥有的所有相同的类方法。
一个 torch.TypedStorage
是一个连续的、一维的数组,其中包含特定 torch.dtype
的元素。它可以指定任何 torch.dtype
,内部数据将得到相应的解释。torch.TypedStorage
包含一个 torch.UntypedStorage
,它将数据存储为无类型的字节数组。
每个跨步 torch.Tensor
都包含一个 torch.TypedStorage
,它存储 torch.Tensor
查看的所有数据。
- class torch.TypedStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- cuda(device=None, non_blocking=False)[source][source]¶
返回此对象在 CUDA 内存中的副本。
如果此对象已在 CUDA 内存中且位于正确的设备上,则不执行复制并返回原始对象。
- property device¶
- classmethod from_file(filename, shared=False, size=0) Storage [source][source]¶
创建一个由内存映射文件支持的 CPU storage。
如果
shared
为True
,则所有进程之间共享内存。所有更改都会写入文件。如果shared
为False
,则 storage 上的更改不会影响文件。size
是存储中的元素数量。如果shared
为False
,则文件必须至少包含size * sizeof(Type)
字节(Type
是存储的类型)。如果shared
为True
,则会在需要时创建文件。- 参数
filename (str) – 要映射的文件名
shared (bool) –
是否共享内存(是否将
MAP_SHARED
或MAP_PRIVATE
传递到底层 mmap(2) 调用)size (int) – storage 中的元素数量
- hpu(device=None, non_blocking=False)[source][source]¶
返回此对象在 HPU 内存中的副本。
如果此对象已在 HPU 内存中且位于正确的设备上,则不执行复制并返回原始对象。
- property is_cuda¶
- property is_hpu¶
- is_pinned(device='cuda')[source][source]¶
确定 CPU TypedStorage 是否已固定在设备上。
- 参数
device (str 或 torch.device) – 要固定内存的设备(默认值:
'cuda'
)。不建议使用此参数,它可能会被弃用。- 返回
一个布尔变量。
- pin_memory(device='cuda')[source][source]¶
将 CPU TypedStorage 复制到锁页内存(如果尚未锁定)。
- 参数
device (str 或 torch.device) – 要固定内存的设备(默认值:
'cuda'
)。不建议使用此参数,它可能会被弃用。- 返回
一个固定内存的 CPU 存储。
- type(dtype=None, non_blocking=False)[source][source]¶
如果未提供 dtype,则返回类型,否则将此对象强制转换为指定的类型。
如果此对象已经是正确的类型,则不执行复制,而是返回原始对象。
- 参数
- 返回类型
Union[_StorageBase, TypedStorage, str]
- untyped()[source][source]¶
返回内部的
torch.UntypedStorage
。
- class torch.DoubleStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.float64[source]¶
- class torch.FloatStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.float32[source]¶
- class torch.HalfStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.float16[source]¶
- class torch.LongStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.int64[source]¶
- class torch.IntStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.int32[source]¶
- class torch.ShortStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.int16[source]¶
- class torch.CharStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.int8[source]¶
- class torch.ByteStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.uint8[source]¶
- class torch.BoolStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.bool[source]¶
- class torch.BFloat16Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.bfloat16[source]¶
- class torch.ComplexDoubleStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.complex128[source]¶
- class torch.ComplexFloatStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.complex64[source]¶
- class torch.QUInt8Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.quint8[source]¶
- class torch.QInt8Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.qint8[source]¶
- class torch.QInt32Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)[source][source]¶
- dtype: torch.dtype = torch.qint32[source]¶