torch.Storage

创建于: 2016年12月30日 | 最后更新于: 2025年04月14日

在 PyTorch 中，一个普通的 tensor 是一个多维数组，它由以下组件定义：

• Storage: tensor 的实际数据，存储为连续的、一维的字节数组。

• dtype: tensor 中元素的ᵢ数据类型，例如 torch.float32 或 torch.int64。

• shape: 一个元组，指示 tensor 在每个维度的ᵢ大小。

• Stride: 在每个维度上从一个元素移动到下一个元素所需的步长。

• Offset: tensor 数据开始ᵢ的 storage 中的起始点。对于新创建的 tensor，这通常为 0。

这些组件共同定义了 tensor 的结构和数据，其中 storage 包含实际数据，其余部分作为元数据。

无类型 Storage API

一个 torch.UntypedStorage 是一个连续的、一维的元素数组。其长度等于 tensor 的字节数。Storage 作为 tensor 的底层数据容器。通常，使用常规构造函数（如 zeros()、zeros_like() 或 new_zeros()）在 PyTorch 中创建的 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 属性，另一个用于其 grad 属性。每个都需要自己的 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 级别：如果存在，应优先使用它们而不是它们的 storage 对应项。

特殊情况

我们提到，具有非 None grad 属性的 tensor 实际上包含两个数据片段。在这种情况下，untyped_storage() 将返回 data 属性的 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 用于形状推理，不包含实际数据。

• Fake Tensors: PyTorch 编译器使用的另一个内部工具是 FakeTensor，它基于类似的想法。

Tensor 子类或类 tensor 对象也可能显示不寻常的行为。总的来说，我们不认为有许多用例需要操作 Storage 级别！

class torch.UntypedStorage(*args, **kwargs)

bfloat16()

将此 storage 转换为 bfloat16 类型。

bool()

将此 storage 转换为 bool 类型。

byte()

将此 storage 转换为 byte 类型。

byteswap(dtype)

交换底层数据的字节。

char()

将此 storage 转换为 char 类型。

clone()

返回此 storage 的副本。

complex_double()

将此 storage 转换为 complex double 类型。

complex_float()

将此 storage 转换为 complex float 类型。

copy_()

cpu()

如果此 storage 尚未在 CPU 上，则返回其 CPU 副本。

cuda(device=None, non_blocking=False)

返回此对象在 CUDA 内存中的副本。

如果此对象已在 CUDA 内存中并且位于正确的设备上，则不执行复制，并返回原始对象。

参数

• device (int) – 目标 GPU ID。默认为当前设备。

• non_blocking (bool) – 如果设置为 True 且源在固定内存中，则复制将与主机异步进行。否则，该参数无效。

返回类型

Union[_StorageBase, TypedStorage]

data_ptr()

device: device

double()

将此 storage 转换为 double 类型。

element_size()

property filename: Optional[str]

返回与此 storage 关联的文件名。

如果 storage 在 CPU 上并且是通过 shared 为 True 的 from_file() 创建的，则文件名将是字符串。否则，该属性为 None。

fill_()

float()

将此 storage 转换为 float 类型。

float8_e4m3fn()

将此 storage 转换为 float8_e4m3fn 类型

float8_e4m3fnuz()

将此 storage 转换为 float8_e4m3fnuz 类型

float8_e5m2()

将此 storage 转换为 float8_e5m2 类型

float8_e5m2fnuz()

将此 storage 转换为 float8_e5m2fnuz 类型

static from_buffer()

static from_file(filename, shared=False, nbytes=0) → Storage

创建一个由内存映射文件支持的 CPU storage。

如果 shared 为 True，则所有进程之间共享内存。所有更改都会写入文件。如果 shared 为 False，则对 storage 的更改不会影响文件。

nbytes 是 storage 的字节数。如果 shared 为 False，则文件必须包含至少 nbytes 字节。如果 shared 为 True，则在需要时会创建该文件。（注意，对于 UntypedStorage，此参数与 TypedStorage.from_file 的参数不同）

参数

• filename (str) – 要映射的文件名

• shared (bool) – 是否共享内存（是否将 MAP_SHARED 或 MAP_PRIVATE 传递给底层 mmap(2) 调用）

• nbytes (int) – storage 的字节数

get_device()

返回类型

int

half()

将此 storage 转换为 half 类型。

hpu(device=None, non_blocking=False)

将此对象复制到 HPU 内存中。

如果此对象已在 HPU 内存中并且位于正确的设备上，则不执行复制，并返回原始对象。

参数

• device (int) – 目标 HPU ID。默认为当前设备。

• non_blocking (bool) – 如果设置为 True 且源在固定内存中，则复制将与主机异步进行。否则，该参数无效。

返回类型

Union[_StorageBase, TypedStorage]

int()

将此 storage 转换为 int 类型。

property is_cuda

property is_hpu

is_pinned(device='cuda')

确定 CPU storage 是否已在 device 上固定。

参数

device (str 或 torch.device) – 要固定内存的设备（默认值：'cuda'）。此参数不推荐使用，且可能会被弃用。

返回

一个布尔变量。

is_shared()

is_sparse: bool = False

is_sparse_csr: bool = False

long()

将此 storage 转换为 long 类型。

mps()

如果此 storage 尚未在 MPS 上，则返回其 MPS 副本。

nbytes()

new()

pin_memory(device='cuda')

如果 CPU storage 尚未固定，则将其复制到固定内存。

参数

device (str 或 torch.device) – 要固定内存的设备（默认值：'cuda'）。此参数不推荐使用，且可能会被弃用。

返回

一个固定的 CPU storage。

resizable()

resize_()

share_memory_(*args, **kwargs)

将 storage 移至共享内存。

对于已在共享内存中的 storage 和 CUDA storage（无需移动即可跨进程共享），此操作无效。共享内存中的 storage 不能调整大小。

请注意，为了缓解 此类 问题，从多个线程调用此函数在同一对象上是线程安全的。但是，在没有适当同步的情况下调用 self 上的任何其他函数则不是线程安全的。有关更多详细信息，请参阅 多进程最佳实践。

注意

当共享内存中 storage 的所有引用都被删除时，关联的共享内存对象也将被删除。PyTorch 有一个特殊的清理过程，以确保即使当前进程意外退出，也能发生这种情况。

值得注意的是 share_memory_() 和 shared = True 的 from_file() 之间的区别。

1. share_memory_ 使用 shm_open(3) 创建一个 POSIX 共享内存对象，而 from_file() 使用 open(2) 打开用户传递的文件名。

2. 两者都使用 MAP_SHARED 标志的 mmap(2) 调用 将文件/对象映射到当前虚拟地址空间。

3. share_memory_ 在映射对象后调用 shm_unlink(3)，以确保在没有进程打开该对象时释放共享内存对象。torch.from_file(shared=True) 不会取消链接该文件。该文件是持久的，将保留直到用户删除。

返回

self

short()

将此 storage 转换为 short 类型。

size()

返回类型

int

to(*, device, non_blocking=False)

tolist()

返回一个包含此 storage 中元素的列表。

type(dtype=None, non_blocking=False)

返回类型

Union[_StorageBase, TypedStorage]

untyped()

旧式类型 Storage

警告

出于历史原因，PyTorch 以前使用类型 Storage 类，这些类现在已弃用，应避免使用。以下详细介绍了此 API，以防您遇到它，尽管其使用强烈不推荐。除了 torch.UntypedStorage 之外，所有 Storage 类都将在未来被移除，并且 torch.UntypedStorage 将在所有情况下使用。

torch.Storage 是与默认数据类型（torch.get_default_dtype()）对应的 Storage 类的别名。例如，如果默认数据类型是 torch.float，则 torch.Storage 解析为 torch.FloatStorage。

像 torch.FloatStorage、torch.IntStorage 等 torch.<type>Storage 和 torch.cuda.<type>Storage 类实际上从不实例化。调用它们的构造函数会创建一个具有适当 torch.dtype 和 torch.device 的 torch.TypedStorage。torch.<type>Storage 类具有 torch.TypedStorage 的所有类方法。

一个 torch.TypedStorage 是一个连续的、一维的、特定 torch.dtype 的元素数组。它可以被赋予任何 torch.dtype，并且内部数据将被正确解释。torch.TypedStorage 包含一个 torch.UntypedStorage，它将数据存储为字节的无类型数组。

每个 strided torch.Tensor 都包含一个 torch.TypedStorage，它存储 torch.Tensor 视图的所有数据。

class torch.TypedStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

bfloat16()

将此 storage 转换为 bfloat16 类型。

bool()

将此 storage 转换为 bool 类型。

byte()

将此 storage 转换为 byte 类型。

char()

将此 storage 转换为 char 类型。

clone()

返回此 storage 的副本。

complex_double()

将此 storage 转换为 complex double 类型。

complex_float()

将此 storage 转换为 complex float 类型。

copy_(source, non_blocking=None)

cpu()

如果此 storage 尚未在 CPU 上，则返回其 CPU 副本。

cuda(device=None, non_blocking=False)

返回此对象在 CUDA 内存中的副本。

如果此对象已在 CUDA 内存中并且位于正确的设备上，则不执行复制，并返回原始对象。

参数

• device (int) – 目标 GPU ID。默认为当前设备。

• non_blocking (bool) – 如果设置为 True 且源在固定内存中，则复制将与主机异步进行。否则，该参数无效。

返回类型

自我

data_ptr()

property device

double()

将此 storage 转换为 double 类型。

dtype: dtype

element_size()

property filename: Optional[str]

如果 storage 是从文件内存映射创建的，则返回与此 storage 关联的文件名；如果 storage 不是通过内存映射文件创建的，则返回 None。

fill_(value)

float()

将此 storage 转换为 float 类型。

float8_e4m3fn()

将此 storage 转换为 float8_e4m3fn 类型

float8_e4m3fnuz()

将此 storage 转换为 float8_e4m3fnuz 类型

float8_e5m2()

将此 storage 转换为 float8_e5m2 类型

float8_e5m2fnuz()

将此 storage 转换为 float8_e5m2fnuz 类型

classmethod from_buffer(*args, **kwargs)

classmethod from_file(filename, shared=False, size=0) → Storage

创建一个由内存映射文件支持的 CPU storage。

如果 shared 为 True，则所有进程之间共享内存。所有更改都会写入文件。如果 shared 为 False，则对 storage 的更改不会影响文件。

size 是 storage 中的元素数量。如果 shared 为 False，则文件必须包含至少 size * sizeof(Type) 字节（Type 是 storage 的类型）。如果 shared 为 True，则在需要时会创建该文件。

参数

• filename (str) – 要映射的文件名

• shared (bool) –

  是否共享内存（是否将 MAP_SHARED 或 MAP_PRIVATE 传递给底层 mmap(2) 调用）

• size (int) – storage 中的元素数量

get_device()

返回类型

int

half()

将此 storage 转换为 half 类型。

hpu(device=None, non_blocking=False)

将此对象复制到 HPU 内存中。

如果此对象已在 HPU 内存中并且位于正确的设备上，则不执行复制，并返回原始对象。

参数

• device (int) – 目标 HPU ID。默认为当前设备。

• non_blocking (bool) – 如果设置为 True 且源在固定内存中，则复制将与主机异步进行。否则，该参数无效。

返回类型

自我

int()

将此 storage 转换为 int 类型。

property is_cuda

property is_hpu

is_pinned(device='cuda')

确定 CPU TypedStorage 是否已在 device 上固定。

参数

device (str 或 torch.device) – 要固定内存的设备（默认值：'cuda'）。此参数不推荐使用，且可能会被弃用。

返回

一个布尔变量。

is_shared()

is_sparse: bool = False

long()

将此 storage 转换为 long 类型。

nbytes()

pickle_storage_type()

pin_memory(device='cuda')

如果 CPU TypedStorage 尚未固定，则将其复制到固定内存。

参数

device (str 或 torch.device) – 要固定内存的设备（默认值：'cuda'）。此参数不推荐使用，且可能会被弃用。

返回

一个固定的 CPU storage。

resizable()

resize_(size)

share_memory_()

参见 torch.UntypedStorage.share_memory_()

short()

将此 storage 转换为 short 类型。

size()

to(*, device, non_blocking=False)

将此对象复制到 device 内存中。

如果此对象已在正确的设备上，则不执行复制，并返回原始对象。

参数

• device (int) – 目标设备。

• non_blocking (bool) – 如果设置为 True 且源在固定内存中，则复制将与主机异步进行。否则，该参数无效。

返回类型

自我

tolist()

返回一个包含此 storage 中元素的列表。

type(dtype=None, non_blocking=False)

如果未提供 dtype，则返回类型，否则将此对象转换为指定类型。

如果已是正确的类型，则不执行复制，并返回原始对象。

参数

• dtype (type 或 string) – 所需类型

• non_blocking (bool) – 如果为 True 且源在固定内存中，目标在 GPU 上（或反之），则复制将与主机异步执行。否则，该参数无效。

• **kwargs – 为了兼容性，可能包含 async 键而不是 non_blocking 参数。async 参数已弃用。

返回类型

Union[_StorageBase, TypedStorage, str]

untyped()

返回内部 torch.UntypedStorage。

class torch.DoubleStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.float64

class torch.FloatStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.float32

class torch.HalfStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.float16

class torch.LongStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.int64

class torch.IntStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.int32

class torch.ShortStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.int16

class torch.CharStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.int8

class torch.ByteStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.uint8

class torch.BoolStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.bool

class torch.BFloat16Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.bfloat16

class torch.ComplexDoubleStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.complex128

class torch.ComplexFloatStorage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.complex64

class torch.QUInt8Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.quint8

class torch.QInt8Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.qint8

class torch.QInt32Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.qint32

class torch.QUInt4x2Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.quint4x2

class torch.QUInt2x4Storage(*args, wrap_storage=None, dtype=None, device=None, _internal=False)

dtype: torch.dtype = torch.quint2x4