PyTorch torch.unbind、torch.split 和 torch.chunk函数介绍

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 则保持原始维度，只是该维度上的大小发生变化。

通过这些示例代码和说明，你可以根据具体需求选择合适的函数来分割多维张量。