PyTorch torch存儲(chǔ)

2025-07-02 11:35 更新

PyTorch 存儲(chǔ)管理詳解

一、什么是 torch.Storage

torch.Storage 是 PyTorch 中用于表示單個(gè)數(shù)據(jù)類型的連續(xù)一維數(shù)組的類。它是 torch.Tensor 的底層數(shù)據(jù)結(jié)構(gòu),每個(gè) torch.Tensor 都對(duì)應(yīng)一個(gè)相同數(shù)據(jù)類型的 torch.Storage。存儲(chǔ)提供了對(duì)張量底層數(shù)據(jù)的直接訪問和操作能力,是內(nèi)存管理的重要組成部分。

二、torch.Storage 的基本操作

(一)創(chuàng)建存儲(chǔ)

您可以使用多種方式創(chuàng)建存儲(chǔ),包括直接構(gòu)造、從文件加載或從緩沖區(qū)創(chuàng)建。

import torch


## 創(chuàng)建一個(gè) FloatStorage
float_storage = torch.FloatStorage(5)  # 創(chuàng)建一個(gè)大小為5的FloatStorage


## 從文件加載存儲(chǔ)
storage_from_file = torch.FloatStorage.from_file("storage.bin", shared=False, size=10)


## 從緩沖區(qū)創(chuàng)建存儲(chǔ)
buffer = bytearray([1, 2, 3, 4])
storage_from_buffer = torch.FloatStorage.from_buffer(buffer, 'big')

(二)類型轉(zhuǎn)換

torch.Storage 支持多種類型轉(zhuǎn)換方法,可以將存儲(chǔ)轉(zhuǎn)換為不同的數(shù)據(jù)類型。

## 將存儲(chǔ)轉(zhuǎn)換為其他類型
char_storage = float_storage.char()
int_storage = float_storage.int()
long_storage = float_storage.long()

(三)數(shù)據(jù)復(fù)制

您可以將存儲(chǔ)復(fù)制到 CPU 或 GPU 上,或者在存儲(chǔ)之間進(jìn)行數(shù)據(jù)復(fù)制。

## 將存儲(chǔ)復(fù)制到CPU
cpu_storage = float_storage.cpu()


## 將存儲(chǔ)復(fù)制到GPU
cuda_storage = float_storage.cuda()


## 復(fù)制存儲(chǔ)數(shù)據(jù)
copied_storage = float_storage.clone()

(四)內(nèi)存共享與固定

torch.Storage 提供了內(nèi)存共享和固定的功能,方便在多進(jìn)程或異步操作中使用。

## 將存儲(chǔ)移動(dòng)到共享內(nèi)存
shared_storage = float_storage.share_memory_()


## 將存儲(chǔ)復(fù)制到固定內(nèi)存
pinned_storage = float_storage.pin_memory()

(五)其他操作

torch.Storage 還支持其他一些操作,如獲取存儲(chǔ)的大小、數(shù)據(jù)類型、元素大小等。

## 獲取存儲(chǔ)的大小
print(float_storage.size())


## 獲取存儲(chǔ)的數(shù)據(jù)類型
print(float_storage.dtype)


## 獲取元素大?。ㄗ止?jié))
print(float_storage.element_size())

三、torch.Storage 的應(yīng)用

(一)高效內(nèi)存管理

在處理大規(guī)模數(shù)據(jù)或長(zhǎng)時(shí)間運(yùn)行的程序時(shí),合理使用存儲(chǔ)可以幫助您更好地管理內(nèi)存,避免內(nèi)存泄漏等問題。

## 示例:創(chuàng)建一個(gè)較大的存儲(chǔ)并手動(dòng)釋放內(nèi)存
import torch


## 創(chuàng)建一個(gè)較大的存儲(chǔ)
large_storage = torch.FloatStorage(1000000)


## 模擬一些操作
## ...


## 手動(dòng)釋放存儲(chǔ)內(nèi)存(在實(shí)際使用中,Python的垃圾回收機(jī)制通常會(huì)自動(dòng)處理)
del large_storage

(二)數(shù)據(jù)持久化

您可以將存儲(chǔ)保存到文件中,以便后續(xù)加載和使用。

## 將存儲(chǔ)保存到文件
storage_to_save = torch.FloatStorage(10)
## 填充存儲(chǔ)數(shù)據(jù)
for i in range(10):
    storage_to_save[i] = i * 1.0


## 保存到文件
storage_to_save.to_file("my_storage.bin")


## 從文件加載存儲(chǔ)
loaded_storage = torch.FloatStorage.from_file("my_storage.bin", shared=False, size=10)
print(loaded_storage.tolist())

(三)多進(jìn)程共享

在多進(jìn)程環(huán)境中,您可以使用共享存儲(chǔ)來實(shí)現(xiàn)進(jìn)程間的數(shù)據(jù)共享。

from multiprocessing import Process
import torch


## 創(chuàng)建共享存儲(chǔ)
shared_storage = torch.FloatStorage.from_file("shared_storage.bin", shared=True, size=5)


def modify_storage(storage):
    # 在子進(jìn)程中修改共享存儲(chǔ)
    for i in range(5):
        storage[i] = i * 2.0


## 創(chuàng)建子進(jìn)程
p = Process(target=modify_storage, args=(shared_storage,))
p.start()
p.join()


## 查看修改后的存儲(chǔ)
print(shared_storage.tolist())

四、總結(jié)

通過本教程,我們?cè)敿?xì)了解了 PyTorch 中的 torch.Storage 類及其相關(guān)操作。torch.Storage 作為 torch.Tensor 的底層數(shù)據(jù)結(jié)構(gòu),在內(nèi)存管理和數(shù)據(jù)操作中扮演著關(guān)鍵角色。掌握存儲(chǔ)的使用方法,可以幫助我們更高效地管理和操作張量數(shù)據(jù),提升程序的性能和穩(wěn)定性。

標(biāo)題:PyTorch 存儲(chǔ)管理詳解:從基礎(chǔ)到實(shí)踐 | 編程獅教程

描述:

關(guān)鍵詞:PyTorch 存儲(chǔ)管理, torch.Storage 應(yīng)用, 內(nèi)存管理技巧, 編程獅 PyTorch 教程, W3Cschool 深度學(xué)習(xí), 存儲(chǔ)類型轉(zhuǎn)換, 數(shù)據(jù)持久化

以上內(nèi)容是否對(duì)您有幫助:
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號(hào)
微信公眾號(hào)

編程獅公眾號(hào)