pytorch中 torch.unbind、torch.split 和 torch.chunk等函数可用于张量的拆分操作。
1. torch.unbind
功能说明：
torch.unbind 沿指定的维度将张量“解包”为多个张量，返回一个元组。解包后被操作的那个维度会消失，每个输出张量的维度数会比原张量少 1。
函数签名：
torch.unbind(input, dim=0)
input: 待解包的张量。dim: 指定沿哪个维度解包。
多维张量示例：
import torch# 创建一个形状为 (2, 3, 4) 的张量
x = torch.arange(24).reshape(2, 3, 4)
print("原始张量 x:\n", x)
print("x.shape:", x.shape) # torch.Size([2, 3, 4])# 沿第 0 维解包
slices0 = torch.unbind(x, dim=0)
print("\n沿 dim=0 解包：")
for i, t in enumerate(slices0):print(f"slice {i} shape: {t.shape}")# 每个张量形状为 (3, 4)# 沿第 1 维解包
slices1 = torch.unbind(x, dim=1)
print("\n沿 dim=1 解包：")
for i, t in enumerate(slices1):print(f"slice {i} shape: {t.shape}")# 每个张量形状为 (2, 4)
2. torch.split
功能说明：
torch.split 根据给定的大小或尺寸列表，将张量沿指定维度切分成若干块。
- 如果传入一个整数，则每块的大小为该整数，最后一块可能会小于这个整数。
- 如果传入一个尺寸列表，则按列表中指定的尺寸进行切分。
函数签名：
torch.split(tensor, split_size_or_sections, dim=0)
tensor: 待分割的张量。split_size_or_sections: 整数或尺寸列表，指定每块的大小。dim: 指定沿哪个维度进行切分。
多维张量示例：
import torch# 创建一个形状为 (2, 5, 4) 的张量
x = torch.arange(40).reshape(2, 5, 4)
print("原始张量 x:\n", x)
print("x.shape:", x.shape) # torch.Size([2, 5, 4])# 按照固定大小进行切分：沿第 1 维，每块大小为 2
splits_fixed = torch.split(x, 2, dim=1)
print("\n沿 dim=1 按固定大小 2 切分：")
for i, t in enumerate(splits_fixed):print(f"chunk {i} shape: {t.shape}")# 输出块的形状可能为 (2, 2, 4), (2, 2, 4) 和最后一块 (2, 1, 4)# 按照指定尺寸列表进行切分：沿第 1 维，分块尺寸为 [1, 2, 2]
splits_list = torch.split(x, [1, 2, 2], dim=1)
print("\n沿 dim=1 按尺寸列表 [1, 2, 2] 切分：")
for i, t in enumerate(splits_list):print(f"chunk {i} shape: {t.shape}")# 分别输出形状 (2, 1, 4), (2, 2, 4), (2, 2, 4)
3. torch.chunk
功能说明：
torch.chunk 将张量沿指定维度平均分成指定数量的块。如果张量在该维度上的长度不能被块数整除，则前面的块会比后面块多一个元素（块的尺寸差别最多为 1）。
函数签名：
torch.chunk(tensor, chunks, dim=0)
tensor: 待分割的张量。chunks: 指定分成几块。dim: 指定沿哪个维度进行分块。
多维张量示例：
对比总结
| 函数 | 分割方式 | 返回结果形式 | 适用场景 |
|---|---|---|---|
| torch.unbind | 沿指定维度将张量完全解包，每个输出不含该维度 | 元组，输出张量数 = 该维度的长度 | 需要逐个处理某一维度上的切片，且希望移除该维度时使用。 |
| torch.split | 按照指定大小或尺寸列表切分张量 | 元组或列表 | 需要按固定大小或自定义尺寸列表切分张量，最后一块可能不均匀。 |
| torch.chunk | 将张量均匀分成指定数量的块 | 元组或列表 | 希望将张量平均分成若干块，块数固定，自动处理无法整除的情况。 |
注意：
- 当处理多维张量时，选择沿哪一维进行分割非常重要；
torch.unbind会移除分割的那个维度，而torch.split和torch.chunk则保持原始维度，只是该维度上的大小发生变化。
通过这些示例代码和说明，你可以根据具体需求选择合适的函数来分割多维张量。
