pytorch api文档:torch.cat()在已有维度上拼接张量(维度数不变)

介绍

torch.cat() 函数,是张量拼接的核心工具,是在已有维度上拼接张量(维度数不变),而 stack 是新增维度堆叠(维度数+1)。

作用

将一系列张量在指定的已有维度(dim)上进行拼接,拼接后张量的维度数与原张量一致,仅指定维度的大小变为所有输入张量在该维度上的大小之和。

语法

torch.cat(tensors, dim=0, out=None) → Tensor

参数:

  • tensors:必选,要拼接的张量序列(列表/元组形式),所有张量除了 dim 维度外,其他维度的形状必须完全相同;
  • dim:可选,指定拼接的维度,默认 0,取值范围是 [-原维度数, 原维度数-1](如2维张量,dim 可取值 -2, -1, 0, 1);
  • out:可选,输出张量,一般无需指定。

与 torch.stack 区别

  • torch.cat:维度数不变,仅指定维度的长度累加;
  • torch.stack:维度数+1,新增维度长度等于张量数量。

例子

先演示不同 dim 的拼接效果,再对比 cat 和 stack 的核心差异:

import torch

# 准备2个张量(满足“除拼接维度外,其他维度形状相同”)
# 张量1:shape [2, 3]
a = torch.tensor([[1, 2, 3], 
                  [4, 5, 6]])
# 张量2:shape [2, 3](和a所有维度都相同,可在任意维度拼接)
b = torch.tensor([[7, 8, 9], 
                  [10, 11, 12]])

示例1:dim=0拼接(行维度)

cat_dim0 = torch.cat([a, b], dim=0)
print("dim=0拼接结果形状:", cat_dim0.shape)  # 输出:torch.Size([4, 3])
print("dim=0拼接结果:")
print(cat_dim0)
# 输出:
# tensor([[ 1,  2,  3],
#         [ 4,  5,  6],
#         [ 7,  8,  9],
#         [10, 11, 12]])

示例2:dim=1拼接(列维度)

cat_dim1 = torch.cat([a, b], dim=1)
print("\ndim=1拼接结果形状:", cat_dim1.shape)  # 输出:torch.Size([2, 6])
print("dim=1拼接结果:")
print(cat_dim1)
# 输出:
# tensor([[ 1,  2,  3,  7,  8,  9],
#         [ 4,  5,  6, 10, 11, 12]])

示例3:非等维度拼接(仅dim维度不同,其他维度相同)

# 张量c:shape [3, 3](dim=0长度3,其他维度和a一致)
c = torch.tensor([[13, 14, 15], 
                  [16, 17, 18],
                  [19, 20, 21]])
cat_dim0_2 = torch.cat([a, c], dim=0)
print("\n不同dim长度拼接结果形状:", cat_dim0_2.shape)  # 输出:torch.Size([5, 3])

对比:torch.stack(新增维度)

stack_dim0 = torch.stack([a, b], dim=0)
print("\ntorch.stack dim=0堆叠结果形状:", stack_dim0.shape)  # 输出:torch.Size([2, 2, 3])

使用场景

torch.cat 是实际开发中最常用的张量拼接工具,典型场景包括:

批量拼接数据

将多个小批次的张量拼接成大批次,或补充数据到已有批次:

# 批次1:batch_size=2, feature_dim=5
batch1 = torch.randn(2, 5)
# 批次2:batch_size=3, feature_dim=5
batch2 = torch.randn(3, 5)
# 拼接成batch_size=5的大批次(dim=0)
total_batch = torch.cat([batch1, batch2], dim=0)
print(total_batch.shape)  # 输出:torch.Size([5, 5])

多特征拼接

在深度学习中,将不同分支的特征张量在通道/维度上拼接(如CNN的多尺度特征融合):

# 特征1:shape [1, 16, 32, 32](batch=1, channel=16, h=32, w=32)
feat1 = torch.randn(1, 16, 32, 32)
# 特征2:shape [1, 32, 32, 32](channel=32,其他维度和feat1一致)
feat2 = torch.randn(1, 32, 32, 32)
# 在通道维度(dim=1)拼接
fusion_feat = torch.cat([feat1, feat2], dim=1)
print(fusion_feat.shape)  # 输出:torch.Size([1, 48, 32, 32])

序列数据拼接

处理NLP中的序列数据,拼接不同句子/段落的张量:

# 句子1:seq_len=10, hidden_dim=768
sent1 = torch.randn(10, 768)
# 句子2:seq_len=15, hidden_dim=768
sent2 = torch.randn(15, 768)
# 拼接成完整文本序列(dim=0)
text_seq = torch.cat([sent1, sent2], dim=0)
print(text_seq.shape)  # 输出:torch.Size([25, 768])

注意事项

  1. 形状匹配要求:除了 dim 指定的维度外,其他所有维度的形状必须完全一致,否则会报错(如不能拼接 [2,3] 和 [2,4] 在 dim=1,但可拼接在 dim=0);
  2. dim的合法范围:dim 不能超过张量的维度数-1(如2维张量最大 dim=1),负数表示倒数维度(dim=-1 等价于最后一维);
  3. 内存效率:torch.cat 不会创建新的内存副本(逻辑拼接),效率高于手动循环拼接;
  4. 与stack的选择:
    • 想在已有维度扩展长度 → 用 cat;
    • 想新增维度组织多个同形状张量 → 用 stack。

总结

  1. torch.cat() 核心是在已有维度拼接张量,维度数不变,仅指定维度的长度累加;
  2. 拼接要求:除 dim 维度外,其他维度形状必须完全相同;
  3. 常用场景包括批量数据拼接、多特征融合、序列数据拼接,是PyTorch中最常用的张量拼接工具,需注意与 torch.stack 的区别。

原文出处:/show_1GW2d1AO7xXt.html