PyTorch张量操作指南：cat、stack、split与chunk的实战拆解

本文深入探讨PyTorch中用于调整张量结构的四个核心函数——torch.cat、torch.stack、torch.split和torch.chunk。通过实际应用场景分析和代码演示，帮助读者掌握它们的功能差异及适用条件，提升模型开发的灵活性与效率。

在深度学习实践中，张量的维度变换是数据处理和模型构建的基础技能。无论是多模态数据的融合（如图像与文本），还是批处理数据的拆分重组，合理运用张量操作函数可显著优化计算流程。PyTorch提供的cat、stack、split和chunk正是解决此类问题的利器。以下将逐一解析其原理与应用。

一、torch.cat: 沿指定维度拼接张量

功能描述

torch.cat（concatenate）沿已有的某一维度连接多个形状兼容的张量，生成更高维度的单一张量。要求除拼接维度外，其余维度的大小必须完全一致。

示例代码

import torch

a = torch.tensor([[1, 2], [3, 4]])  # 形状 (2, 2)
b = torch.tensor([[5, 6], [7, 8]])

# 在第0维拼接（垂直方向）
c = torch.cat([a, b], dim=0)  
print(c)
# 输出:
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])

# 在第1维拼接（水平方向）
d = torch.cat([a, b], dim=1)  
print(d)
# 输出:
# tensor([[1, 2, 5, 6],
#         [3, 4, 7, 8]])

二、torch.stack: 创建新维度堆叠张量

功能描述

torch.stack会将输入张量沿新创建的维度进行堆叠，所有参与堆叠的张量必须具有完全相同的形状。输出张量的维度比原张量多一维。

示例代码

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])

# 沿第0维堆叠，生成二维张量
c = torch.stack([a, b], dim=0)  
print(c.shape)  # torch.Size([2, 3])
print(c)
# 输出:
# tensor([[1, 2, 3],
#         [4, 5, 6]])

# 沿第1维堆叠，生成二维张量
d = torch.stack([a, b], dim=1)  
print(d.shape)  # torch.Size([3, 2])
print(d)
# 输出:
# tensor([[1, 4],
#         [2, 5],
#         [3, 6]])

三、torch.split: 按尺寸分割张量

功能描述

torch.split根据指定的尺寸将输入张量分割为多个子张量。支持两种参数形式：

• 整数列表：每个元素表示对应分片的长度
• 整数N：等分为N个子张量（需总长度可被整除）

示例代码

a = torch.arange(9)  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])

# 按列表尺寸分割 [2,3,4]
parts = torch.split(a, [2, 3, 4], dim=0)
for part in parts:
    print(part)

'''
输出:
tensor([0, 1])
tensor([2, 3, 4])
tensor([5, 6, 7, 8])
'''

# 平均分割为3份
chunks = torch.split(a, 3, dim=0)
print([c.shape for c in chunks])  # [torch.Size([3]), torch.Size([3]), torch.Size([3])]

四、torch.chunk: 按数量均分张量

功能描述

torch.chunk将输入张量沿指定维度均匀划分为N份。若无法整除，剩余元素分配到前面的分片中。

示例代码

a = torch.arange(10)  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

# 分成3份，默认在第0维操作
chunks = torch.chunk(a, chunks=3, dim=0)
for i, chunk in enumerate(chunks):
    print(f"Chunk {i}: {chunk}")

'''
输出:
Chunk 0: tensor([0, 1, 2, 3])
Chunk 1: tensor([4, 5, 6])
Chunk 2: tensor([7, 8, 9])
'''

# 在第1维分割二维张量
b = a.reshape(2,5)
chunks = torch.chunk(b, chunks=2, dim=1)
print(chunks[0].shape)  # torch.Size([2, 2])
print(chunks[1].shape)  # torch.Size([2, 3])

综合示例

以下是结合图像数据的完整操作示例，模拟图像预处理流程中的张量操作场景：

综合示例：图像数据的分割与合并处理

场景设定

假设我们有一批RGB图像数据（尺寸为 3×256×256），需要完成以下操作：

1. 将图像拆分为RGB三个通道
2. 对每个通道进行独立归一化
3. 合并处理后的通道
4. 将多张图像堆叠成批次
5. 分割批次为训练/验证集

代码实现

import torch
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt

# 1. 加载示例图像 (H, W, C) -> 转换为 (C, H, W)
image = Image.open('cat.jpg').convert('RGB')
image = transforms.ToTensor()(image)  # shape: torch.Size([3, 256, 256])

# 2. 使用split分离RGB通道
r_channel, g_channel, b_channel = torch.split(image, split_size_or_sections=1, dim=0)

''' 可视化原始通道
plt.figure(figsize=(12,4))
plt.subplot(131), plt.imshow(r_channel.squeeze().numpy(), cmap='Reds'), plt.title('Red')
plt.subplot(132), plt.imshow(g_channel.squeeze().numpy(), cmap='Greens'), plt.title('Green')
plt.subplot(133), plt.imshow(b_channel.squeeze().numpy(), cmap='Blues'), plt.title('Blue')
plt.show()
'''

# 3. 对每个通道进行归一化（示例操作）
def normalize(tensor):
    return (tensor - tensor.mean()) / tensor.std()

r_norm = normalize(r_channel)
g_norm = normalize(g_channel)
b_norm = normalize(b_channel)

# 4. 使用cat合并处理后的通道
normalized_img = torch.cat([r_norm, g_norm, b_norm], dim=0)
'''观察归一化效果
plt.imshow(normalized_img.permute(1,2,0))
plt.title('Normalized Image')
plt.show()
'''

# 5. 创建模拟图像批次 (假设有4张相同图像)
batch_images = torch.stack([image]*4, dim=0)  # shape: (4, 3, 256, 256)

# 6. 使用chunk分割批次为训练集/验证集
train_set, val_set = torch.chunk(batch_images, chunks=2, dim=0)
print(f"Train set size: {train_set.shape}")  # torch.Size([2, 3, 256, 256])
print(f"Val set size: {val_set.shape}")      # torch.Size([2, 3, 256, 256])

关键操作解析

扩展应用建议

1. 数据增强：对split后的通道进行不同变换（如仅对R通道做对比度调整）
2. 模型输入：stack后的批次可直接输入CNN网络
3. 分布式训练：利用chunk将数据分布到多个GPU处理
4. 特征可视化：通过split提取中间层特征图的单个通道进行分析

通过这个完整的图像处理流程示例，可以清晰看到：

• split+cat 组合常用于特征处理管道
• stack+chunk 组合是构建批处理系统的关键工具
• 这些操作在保持计算效率的同时提供了灵活的数据控制能力

总结与对比

应用建议：

• 当需要合并同类数据且保留原始维度时用cat；
• 若需扩展维度以表示批次或通道时用stack；
• 对序列数据分段处理优先考虑split；
• 均匀划分特征图或张量时选择chunk。

掌握这些工具后，您将能更灵活地操控张量维度，适应复杂模型的构建需求！