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