当前位置: 首页 > article >正文

PyTorch 张量的分块处理介绍

分块处理是将大型张量分解成较小的块,以便更高效地进行计算,减少内存占用,特别适用于处理超大张量的场景(如深度学习中的大批量数据或大型模型训练)。

PyTorch 提供了多种方法来分块张量,包括 chunksplitunfold 等。以下是这些方法的介绍及其适用场景。


1. 使用 torch.chunk

功能
将张量沿指定维度分成 指定数量的块

语法

torch.chunk(input, chunks, dim=0)
  • input:输入张量。
  • chunks:将张量分成的块数。
  • dim:指定分块的维度。

示例

x = torch.arange(12).reshape(3, 4)  # 一个 3x4 的张量
print(x)

# 沿第 0 维(行)分成 3 块
chunks = torch.chunk(x, chunks=3, dim=0)
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}:\n{chunk}")

2. 使用 torch.split

功能
将张量沿指定维度分成 指定大小的块。每块的大小可以不相同。

语法

torch.split(input, split_size_or_sections, dim=0)
  • split_size_or_sections:可以是整数(每块的大小)或列表(每块的大小分别定义)。
  • dim:指定分块的维度。

示例

x = torch.arange(12).reshape(3, 4)

# 沿第 1 维(列)分成大小为 2 的块
chunks = torch.split(x, 2, dim=1)
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}:\n{chunk}")

# 指定块的大小分别为 1 和 3
chunks = torch.split(x, [1, 3], dim=1)
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}:\n{chunk}")

3. 使用 torch.unfold

功能
从张量中提取 滑动窗口 块,适用于卷积、时间序列等需要处理局部区域的场景。

语法

torch.unfold(input, dimension, size, step)
  • dimension:滑动窗口应用的维度。
  • size:窗口大小。
  • step:滑动步长。

示例

x = torch.arange(1, 10).reshape(1, 1, 3, 3)  # 3x3 张量
print("Original tensor:\n", x)

# 滑动窗口大小为 2x2,步长为 1
unfolded = x.unfold(2, 2, 1).unfold(3, 2, 1)
print("Unfolded tensor shape:", unfolded.shape)
print("Unfolded tensor:\n", unfolded)

注: x.unfold(2, 2, 1)中第一个2为维度,第二个2为窗口大小,1为滑动步长。

4. 分块处理的优化场景

(1) 内存优化
  • 当内存不足以加载整个大张量时,将其分块并逐块处理可以显著减少内存占用。
(2) 并行处理
  • 将张量分块后,可以将各个块分配到不同的设备(如多 GPU)进行并行计算,提高计算效率。
(3) 数据流批处理
  • 在训练深度学习模型时,可将大批次数据分块处理,避免一次性加载超大批次导致显存溢出。

5. 分块与重组

分块后常需要对处理后的块进行重组,以恢复原始张量的形状。PyTorch 提供以下方法:

(1) torch.cat

将多个张量按指定维度连接。

x1 = torch.tensor([[1, 2]])
x2 = torch.tensor([[3, 4]])
result = torch.cat([x1, x2], dim=0)
print(result)  # tensor([[1, 2],
               #         [3, 4]])
(2) torch.stack

将多个张量堆叠到新的维度上。

x1 = torch.tensor([1, 2])
x2 = torch.tensor([3, 4])
result = torch.stack([x1, x2], dim=0)
print(result)  # tensor([[1, 2],
               #         [3, 4]])
(3) torch.reshape

用于恢复分块前的形状。

chunk = torch.arange(12).reshape(3, 4)
flattened = chunk.flatten()
restored = flattened.reshape(3, 4)
print(restored)

6. 总结

  • chunk:按块数分割张量,块大小可能不均。
  • split:按块大小分割张量,可以自定义每块的大小。
  • unfold:生成滑动窗口块,适用于局部区域计算。
  • 分块优化场景:适用于内存优化、并行计算和大批次数据处理。


http://www.kler.cn/a/502428.html

相关文章:

  • 鸿蒙UI开发——带农历的日期滑动选择弹窗
  • 74 mysql having 的实现
  • 数据结构与算法之链表: LeetCode 234. 回文链表 (Ts版)
  • sql server 对 nvarchar 类型的列进行 SUM() 运算
  • Spring Boot 动态表操作服务实现
  • OS1.【Linux】大致介绍和环境搭建
  • Redis高危漏洞-GHSA-whxg-wx83-85p5:用户可能会使用特制的 Lua 脚本来触发堆栈缓冲区溢出
  • uc/os-II 原理及应用(八) 系统裁减以及移植到51单片机上
  • 掌握 Ubuntu 终端 mv 与 rename 命令的高效重命名使用方法
  • STM32-笔记42-实时时钟项目
  • uniapp 抖音小程序 getUserProfile:fail must be invoked by user tap gesture
  • CMake学习笔记(1)
  • 开源免费的下载工具AB Download Manager
  • 中等难度——python实现电子宠物和截图工具
  • 概率输出和独热分割掩码的主要区别:
  • 每日学习30分轻松掌握CursorAI:Cursor基础设置与配置
  • 商用服务器密码机的加密技术与优势
  • Win32汇编学习笔记11.游戏辅助的实现
  • fft分析数据求bode图原理
  • 【SQL】进阶知识 -- 删除表的几种方法(包含表内单个字段的删除方法)