PyTorch 基础学习（11）- 张量内存管理

系列文章：
PyTorch 基础学习（1） - 快速入门
PyTorch 基础学习（2）- 张量 Tensors
PyTorch 基础学习（3） - 张量的数学操作
PyTorch 基础学习（4）- 张量的类型
PyTorch 基础学习（5）- 神经网络
PyTorch 基础学习（6）- 函数API
PyTorch 基础学习（7）- 自动微分
PyTorch 基础学习（8）- 多进程并发
PyTorch 基础学习（9）- 训练优化器
PyTorch 基础学习（10）- Transformer

介绍

在 PyTorch 中，torch.Storage 是一个基础数据结构，用于存储张量（Tensor）背后的数据。Storage 是一个一维的连续数组，每个元素都是相同的数据类型。虽然用户通常直接与张量交互，但 Storage 在幕后支持张量的内存管理。

每个 torch.Tensor 都有一个对应的 Storage 对象，该对象保存了张量的数据。通过 Storage，可以访问和操作张量底层的数据。

重要方法及其作用

以下是 torch.FloatStorage 类中的一些重要方法及其功能：

• byte(): 将存储的数据类型转换为 byte 类型。
• char(): 将存储的数据类型转换为 char 类型。
• clone(): 返回存储的一个副本。
• cpu(): 如果当前存储不在 CPU 上，则返回它在 CPU 上的一个副本。
• cuda(device=None, async=False): 返回此对象在 CUDA 内存中的副本。可以指定目标 GPU 的 ID。async 参数允许在某些情况下异步复制。
• double(): 将存储的数据类型转换为 double 类型。
• float(): 将存储的数据类型转换为 float 类型。
• half(): 将存储的数据类型转换为 half 类型。
• int(): 将存储的数据类型转换为 int 类型。
• long(): 将存储的数据类型转换为 long 类型。
• resize_(): 改变存储的大小。
• tolist(): 将存储中的所有元素转换为 Python 列表。
• type(new_type=None, async=False): 将存储的数据类型转换为指定的类型。

主要使用场景

torch.Storage 主要用于以下场景：

• 内存管理: 当需要手动管理张量的数据时，可以直接使用 Storage。
• 类型转换: 通过存储的类型转换方法，可以方便地在不同的数据类型之间切换。
• 设备切换: 在 CPU 和 GPU 之间迁移数据时，可以使用 cpu() 和 cuda() 方法。
• 数据共享: share_memory_() 方法可以将存储的数据移动到共享内存中，使得多个进程能够访问相同的数据。

应用实例

以下是一个简单的例子，展示了如何使用 torch.Storage 来管理张量的底层数据。

import torch

# 创建一个浮点型张量
tensor = torch.FloatTensor([1.0, 2.0, 3.0, 4.0])

# 获取张量的存储
storage = tensor.storage()

# 查看存储中的元素
print("Storage elements:", storage.tolist())

# 将存储的数据类型转换为 byte
storage.byte()

# 查看转换后的存储类型
print("Converted storage elements (byte):", storage.tolist())

# 将存储移到 GPU（如果可用）
if torch.cuda.is_available():
    storage_cuda = storage.cuda()
    print("Storage moved to CUDA")

# 克隆存储
storage_clone = storage.clone()
print("Cloned storage elements:", storage_clone.tolist())

小结

torch.Storage 是 PyTorch 中的一个底层数据结构，用于管理张量的内存。通过理解 Storage 的基本概念和常用方法，开发者可以更灵活地操作和管理张量的数据，从而在更高级的场景中应用 PyTorch。