企业官网建设流程全解析

1. 为什么张量拆分与拼接是绕不开的基本功

刚接触PyTorch那会儿，我总觉得张量操作里最不起眼的就是torch.split、torch.chunk、torch.cat、torch.stack这几个函数。直到有一次在做一个多分支网络的特征融合模块，输入特征图需要按通道切成四份分别送进不同的卷积分支，处理完再拼回来。我随手写了torch.chunk(x, 4, dim=1)，结果通道数不能被4整除，最后一份的尺寸和其他三份不一样，直接导致后续的卷积层报维度错误。排查了大半天才发现问题出在这个看似简单的拆分操作上。

这件事让我意识到，张量的拆分与拼接虽然API简单，但真正用好需要理解它们各自的行为边界、内存布局的影响、以及在不同维度上操作时的语义差异。这篇文章就是把我这些年在这几个函数上踩过的坑、总结的经验，以及一些容易被忽略的细节，系统地梳理一遍。无论你是刚入门PyTorch的新手，还是已经写过不少模型的老手，相信都能从中找到一些之前没注意到的点。

核心关键词：PyTorch张量拆分、张量拼接、torch.split、torch.chunk、torch.cat、torch.stack、维度操作。

2. 拆分操作：split、chunk与unbind的语义差异

2.1 torch.split：按指定大小精确切分

torch.split是我日常用得最多的拆分函数，它的行为非常直观：给定一个张量和一个分割大小（或分割大小列表），沿着指定维度切开。

import torch x = torch.arange(24).reshape(4, 6) print(x) # tensor([[ 0, 1, 2, 3, 4, 5], # [ 6, 7, 8, 9, 10, 11], # [12, 13, 14, 15, 16, 17], # [18, 19, 20, 21, 22, 23]]) # 沿dim=0，每份大小为2 parts = torch.split(x, 2, dim=0) print(len(parts)) # 2 print(parts[0].shape) # torch.Size([2, 6]) print(parts[1].shape) # torch.Size([2, 6]) # 沿dim=1，每份大小为3 parts = torch.split(x, 3, dim=1) print(len(parts)) # 2 print(parts[0].shape) # torch.Size([4, 3])

这里有一个关键细节：当张量在指定维度上的长度不能被分割大小整除时，最后一份会小于指定的分割大小。比如把长度为7的维度按大小3拆分，会得到3、3、1三份。这个行为在写代码时如果不注意，很容易在后续处理中出问题。

y = torch.arange(7) parts = torch.split(y, 3) print([p.shape[0] for p in parts]) # [3, 3, 1]

torch.split还支持传入一个列表来指定每一份的大小，这在需要不均匀拆分时非常有用：

x = torch.arange(10) parts = torch.split(x, [2, 3, 5]) print([p.shape[0] for p in parts]) # [2, 3, 5]

注意：传入列表时，列表元素之和必须等于该维度的长度，否则会报错。这一点和torch.chunk的行为不同，后面会对比。

2.2 torch.chunk：按份数尽量均匀地切分

torch.chunk的语义是"我想把这个维度切成N份"，而不是"每份多大"。它内部会计算每份的大小，然后调用类似split的逻辑。

x = torch.arange(24).reshape(4, 6) # 沿dim=1切成3份 parts = torch.chunk(x, 3, dim=1) print([p.shape for p in parts]) # [torch.Size([4, 2]), torch.Size([4, 2]), torch.Size([4, 2])] # 沿dim=1切成4份（6不能被4整除） parts = torch.chunk(x, 4, dim=1) print([p.shape for p in parts]) # [torch.Size([4, 2]), torch.Size([4, 2]), torch.Size([4, 1]), torch.Size([4, 1])]

从上面的输出可以看到，当不能整除时，chunk会尽量均匀分配，前面的份数较大，后面的份数较小。具体来说，它计算每份大小为ceil(6/4)=2，然后按这个大小切分，最后得到2、2、1、1。

这里有一个容易踩的坑：chunk返回的份数不一定等于你指定的份数。比如：

x = torch.arange(3) parts = torch.chunk(x, 5) print(len(parts)) # 3，不是5 print([p.shape[0] for p in parts]) # [1, 1, 1]

当维度长度小于指定的份数时，chunk最多只能切出长度个份数（每份大小为1）。这个行为在写通用代码时如果不做检查，可能会导致后续逻辑出错。

2.3 torch.unbind：沿维度完全拆开

torch.unbind是把一个维度彻底拆掉，返回一个元组，每个元素是去掉该维度后的张量。可以理解为split每份大小为1的特例。

x = torch.arange(24).reshape(4, 6) # 沿dim=0拆开 rows = torch.unbind(x, dim=0) print(len(rows)) # 4 print(rows[0].shape) # torch.Size([6]) # 沿dim=1拆开 cols = torch.unbind(x, dim=1) print(len(cols)) # 6 print(cols[0].shape) # torch.Size([4])

unbind的一个典型使用场景是处理序列数据时，需要逐时间步操作。但要注意，unbind返回的是原张量的视图（view），对返回的张量做原地修改会影响原张量。这一点和split、chunk一样，它们返回的都是视图而非拷贝。

2.4 三者的选择依据与常见误用

我个人的选择习惯是：需要精确控制每份大小时用split，只是想把东西分成几堆时用chunk，需要完全拆开时用unbind。在写需要保证输出形状的代码时，尽量用split并传入列表，这样行为完全确定，不会因为输入尺寸变化而出现意外。

还有一个常见的误用：有人会用chunk来做数据集的划分，比如把1000个样本分成8个batch。如果1000不能被8整除，chunk会返回8份但大小不一（125、125、...、125、...），而实际上你可能想要的是每个batch大小一致，最后不足的丢弃或补齐。这种情况下应该用split配合手动计算，或者直接用DataLoader。

3. 拼接操作：cat与stack的本质区别

3.1 torch.cat：沿已有维度连接

torch.cat是在现有维度上把多个张量接起来，不增加新的维度。这是最常用的拼接方式。

a = torch.randn(2, 3) b = torch.randn(2, 4) # 沿dim=1拼接 c = torch.cat([a, b], dim=1) print(c.shape) # torch.Size([2, 7]) # 沿dim=0拼接 d = torch.cat([a, a], dim=0) print(d.shape) # torch.Size([4, 3])

cat的核心约束是：除了拼接维度外，其他所有维度的大小必须完全一致。比如上面沿dim=1拼接时，dim=0的大小必须都是2。如果不一致，会直接报错。

a = torch.randn(2, 3) b = torch.randn(3, 4) try: torch.cat([a, b], dim=1) except RuntimeError as e: print(e) # Sizes of tensors must match except in dimension 1. # Expected size 2 but got size 3 for tensor number 1 in the list.

这个报错信息其实很清晰，它告诉你除了dim=1之外，其他维度期望的大小和实际得到的大小。排查时重点看报错里说的"except in dimension X"之外的维度。

3.2 torch.stack：创建新维度堆叠

torch.stack是在新维度上堆叠张量，要求所有张量的形状完全一致。

a = torch.randn(3, 4) b = torch.randn(3, 4) c = torch.randn(3, 4) # 沿dim=0堆叠 s = torch.stack([a, b, c], dim=0) print(s.shape) # torch.Size([3, 3, 4]) # 沿dim=1堆叠 s = torch.stack([a, b, c], dim=1) print(s.shape) # torch.Size([3, 3, 4]) # 沿dim=2堆叠 s = torch.stack([a, b, c], dim=2) print(s.shape) # torch.Size([3, 4, 3])

stack和cat最容易混淆的地方在于：cat是在已有维度上操作，维度数不变；stack是新增一个维度，维度数加1。我见过不少人在需要堆叠时用了cat，结果形状不对，排查半天。

一个简单的记忆方法：如果你想把N个形状相同的张量组合成一个"批次"，用stack；如果你想把两个张量在某个维度上接起来，用cat。

3.3 拼接时的维度对齐陷阱

cat和stack在指定dim时，负索引的行为需要特别注意。PyTorch支持负数索引，dim=-1表示最后一个维度。但在stack中，dim的取值范围是[-(N+1), N]，其中N是输入张量的维度数。因为stack会新增一个维度，所以可选的dim比输入维度多一个。

a = torch.randn(3, 4) b = torch.randn(3, 4) # stack时dim=-1表示在最后新增维度 s = torch.stack([a, b], dim=-1) print(s.shape) # torch.Size([3, 4, 2]) # stack时dim=-3表示在最前面新增维度 s = torch.stack([a, b], dim=-3) print(s.shape) # torch.Size([2, 3, 4])

另一个容易忽略的点是：cat和stack都不会改变张量的数据类型。如果拼接的张量类型不一致（比如一个是float32，一个是float64），会报错。在混合精度训练中，这一点尤其需要注意，因为模型参数可能是float16，而某些中间计算结果可能是float32。

3.4 性能考量：拼接操作的隐藏成本

cat和stack都会分配新的内存来存储结果，这意味着它们不是原地操作，会带来内存分配和数据拷贝的开销。在循环中频繁拼接张量是一个常见的性能反模式。

# 反模式：循环中不断cat result = torch.tensor([]) for i in range(1000): result = torch.cat([result, torch.randn(1)]) # 推荐：先收集到列表，最后一次性cat parts = [] for i in range(1000): parts.append(torch.randn(1)) result = torch.cat(parts)

上面第一种写法的时间复杂度是O(n²)，因为每次cat都要拷贝之前的所有数据。第二种写法是O(n)，只在最后拷贝一次。在处理大规模数据时，这个差异可能是几秒和几分钟的区别。

提示：如果确实需要在循环中逐步构建张量，可以考虑预先分配一个足够大的张量，然后用索引赋值。但这种方式需要提前知道最终大小，灵活性较差。

4. 维度变换中的拆分拼接组合技

4.1 用view/reshape配合拆分实现维度重排

拆分和拼接经常和view、reshape、permute、transpose配合使用，来实现复杂的维度变换。一个经典的例子是多头注意力中的维度重排。

假设我们有一个形状为(batch, seq_len, d_model)的张量，需要拆成(batch, num_heads, seq_len, d_head)，其中d_model = num_heads * d_head。

batch, seq_len, d_model = 2, 10, 512 num_heads = 8 d_head = d_model // num_heads x = torch.randn(batch, seq_len, d_model) # 方法一：先view再transpose x1 = x.view(batch, seq_len, num_heads, d_head) x1 = x1.transpose(1, 2) # (batch, num_heads, seq_len, d_head) # 方法二：先split再stack parts = torch.split(x, d_head, dim=-1) # num_heads个(batch, seq_len, d_head) x2 = torch.stack(parts, dim=1) # (batch, num_heads, seq_len, d_head) print(torch.equal(x1, x2)) # True

两种方法结果等价，但方法一更高效，因为它只涉及view和transpose，不涉及数据拷贝（transpose返回的是视图）。方法二中的stack会分配新内存并拷贝数据。所以在实际写模型时，优先用view + transpose的组合。

4.2 用cat实现特征融合的常见模式

在特征融合场景中，cat是最常用的操作。比如U-Net中的跳跃连接，就是把编码器的特征和解码器的特征沿通道维度拼接。

# 模拟U-Net中的跳跃连接 encoder_feat = torch.randn(2, 64, 32, 32) # 编码器特征 decoder_feat = torch.randn(2, 128, 32, 32) # 解码器特征 # 沿通道维度拼接 fused = torch.cat([decoder_feat, encoder_feat], dim=1) print(fused.shape) # torch.Size([2, 192, 32, 32])

这里有一个实践中的经验：拼接的顺序会影响后续卷积层的权重学习。虽然理论上先拼谁后拼谁不影响模型的表达能力，但在实际训练中，保持一致的拼接顺序有助于稳定训练。我习惯把"主路"特征放在前面，"旁路"特征放在后面。

4.3 拆分后并行处理再拼接的完整案例

下面是一个完整的例子，模拟多分支卷积块：输入特征按通道分成四组，分别经过不同大小的卷积核，最后拼接回来。

import torch.nn as nn class MultiBranchBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() # 确保通道数能被4整除 assert in_channels % 4 == 0, "in_channels必须能被4整除" branch_channels = in_channels // 4 self.branch1 = nn.Conv2d(branch_channels, out_channels // 4, 1) self.branch2 = nn.Conv2d(branch_channels, out_channels // 4, 3, padding=1) self.branch3 = nn.Conv2d(branch_channels, out_channels // 4, 5, padding=2) self.branch4 = nn.Conv2d(branch_channels, out_channels // 4, 7, padding=3) def forward(self, x): # 沿通道维度分成4份 parts = torch.chunk(x, 4, dim=1) # 分别处理 out1 = self.branch1(parts[0]) out2 = self.branch2(parts[1]) out3 = self.branch3(parts[2]) out4 = self.branch4(parts[3]) # 拼接回来 return torch.cat([out1, out2, out3, out4], dim=1)

这个模式在ResNeXt、Inception等网络中都有应用。关键点是拆分前要确保通道数能被份数整除，否则chunk返回的份数大小不一致，后续的卷积层就无法统一处理了。我在代码里加了一个assert来提前暴露问题，这比等到运行时维度报错要好得多。

4.4 非连续张量上的拆分拼接注意事项

PyTorch中的张量有"连续"和"非连续"之分。transpose、permute、expand等操作返回的是非连续张量（视图），在非连续张量上做view会报错，需要先调用.contiguous()。

x = torch.randn(4, 6) y = x.transpose(0, 1) # 非连续 print(y.is_contiguous()) # False # 在非连续张量上split是可以的 parts = torch.split(y, 2, dim=0) print(parts[0].is_contiguous()) # False # 但如果要对拆分结果做view，需要先contiguous try: parts[0].view(-1) except RuntimeError as e: print(e) # view size is not compatible with input tensor's size and stride... # 正确做法 parts[0].contiguous().view(-1)

cat和stack对非连续张量是兼容的，它们会正确处理步长。但如果你在拼接后要做view，最好在拼接前或拼接后调用.contiguous()，避免运行时错误。

5. 实战中容易踩的坑与排查思路

5.1 维度不匹配报错的快速定位方法

维度不匹配是拆分拼接中最常见的错误。PyTorch的报错信息通常会告诉你哪个维度不匹配，但有时候信息比较长，容易看漏。我的排查习惯是：

1. 先看报错最后一行，通常会给出期望大小和实际大小。
2. 打印所有输入张量的shape，在拼接前加一行print([t.shape for t in tensors])。
3. 检查拼接维度是否正确，特别是用了负数索引时。

# 一个典型的排查辅助函数 def safe_cat(tensors, dim=0): shapes = [t.shape for t in tensors] print(f"拼接前形状: {shapes}, 拼接维度: {dim}") # 检查除dim外其他维度是否一致 for i in range(len(shapes[0])): if i == dim % len(shapes[0]): continue sizes = [s[i] for s in shapes] if len(set(sizes)) != 1: raise ValueError(f"维度{i}大小不一致: {sizes}") return torch.cat(tensors, dim=dim)

这个辅助函数在调试阶段很有用，能快速定位是哪个维度出了问题。

5.2 拆分后形状不确定导致的后续错误

chunk的不均匀拆分是隐蔽性很强的坑。比如你写了一个模块，假设输入通道是64，用chunk(x, 3, dim=1)分成3份，你以为每份是21、21、22或者类似，但实际上ceil(64/3)=22，所以是22、22、20。如果你的后续代码假设每份都是22，就会出错。

x = torch.randn(2, 64, 8, 8) parts = torch.chunk(x, 3, dim=1) print([p.shape[1] for p in parts]) # [22, 22, 20]

避免这个坑的最好方法是：要么确保能整除，要么用split传入明确的列表。

# 明确指定每份大小 parts = torch.split(x, [22, 22, 20], dim=1)

5.3 原地操作与视图共享的内存陷阱

split、chunk、unbind返回的都是视图，和原张量共享内存。如果你修改了拆分结果，原张量也会被修改。

x = torch.arange(12).reshape(3, 4) parts = torch.split(x, 1, dim=0) parts[0].fill_(999) print(x) # tensor([[999, 999, 999, 999], # [ 4, 5, 6, 7], # [ 8, 9, 10, 11]])

这个行为在大多数时候是符合预期的（也是高效的），但如果你不小心在某个地方做了原地修改，可能会影响到其他使用原张量的代码。如果需要独立的副本，显式调用.clone()。

parts = [p.clone() for p in torch.split(x, 1, dim=0)]

5.4 在GPU上的拆分拼接性能实测

在GPU上，拆分和拼接的开销比CPU上更值得关注，因为GPU的内存分配和内核启动都有额外开销。我做了一个简单的测试：

import torch import time device = torch.device('cuda') x = torch.randn(64, 256, 32, 32, device=device) # 测试chunk + cat的开销 torch.cuda.synchronize() start = time.time() for _ in range(1000): parts = torch.chunk(x, 4, dim=1) y = torch.cat(parts, dim=1) torch.cuda.synchronize() print(f"chunk+cat耗时: {time.time() - start:.4f}s") # 对比：直接clone torch.cuda.synchronize() start = time.time() for _ in range(1000): y = x.clone() torch.cuda.synchronize() print(f"clone耗时: {time.time() - start:.4f}s")

实测下来，chunk+cat的耗时大约是clone的2-3倍。虽然单次开销不大，但在训练循环中如果每个step都做多次拆分拼接，累积起来就不可忽略了。优化思路是：能用一个view或reshape搞定的，就不要用拆分拼接。

5.5 自动求导对拆分拼接的支持边界

split、chunk、cat、stack都是可微分的，梯度会正确地反向传播。但有一个细节需要注意：unbind的梯度传播在某些PyTorch版本中可能不如split稳定。我遇到过在复杂计算图中使用unbind后梯度为None的情况，换成split就正常了。虽然这可能是特定版本的bug，但在生产代码中，我倾向于用split替代unbind来规避潜在问题。

另外，cat的梯度会按照拼接时的划分反向传回各个输入张量，这是自动完成的，不需要额外处理。但如果你在拼接后做了原地操作（比如y += something），可能会破坏计算图，导致梯度计算出错。

6. 几个高频场景的代码模板

6.1 序列数据的分段处理

处理长序列时，经常需要把序列切成固定长度的段。下面是一个模板：

def split_sequence(x, seg_len, dim=1): """将序列按seg_len切分，不足的部分用0填充""" seq_len = x.shape[dim] num_segs = (seq_len + seg_len - 1) // seg_len pad_len = num_segs * seg_len - seq_len if pad_len > 0: # 构造padding pad_shape = list(x.shape) pad_shape[dim] = pad_len padding = torch.zeros(pad_shape, dtype=x.dtype, device=x.device) x = torch.cat([x, padding], dim=dim) # 用split切分 segments = torch.split(x, seg_len, dim=dim) return segments

这个模板的关键点是先补齐再切分，保证每段长度一致。如果直接chunk，最后一段可能较短，后续处理时需要额外判断。

6.2 多模型输出的集成拼接

在做模型集成时，经常需要把多个模型的输出拼接起来：

def ensemble_concat(outputs, dim=-1): """将多个模型的输出沿最后一维拼接""" # 确保所有输出形状一致（除了拼接维度） shapes = [o.shape for o in outputs] for i in range(len(shapes[0])): if i == dim % len(shapes[0]): continue assert len(set(s[i] for s in shapes)) == 1, f"维度{i}不一致" return torch.cat(outputs, dim=dim)

6.3 通道注意力的拆分-加权-拼接

通道注意力机制（如SE Block）中，经常需要把特征图拆开、加权、再拼回去：

class ChannelAttention(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.fc = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, h, w = x.shape # 全局平均池化 y = x.mean(dim=[2, 3]) # (b, c) # 计算注意力权重 weights = self.fc(y).view(b, c, 1, 1) # 加权 return x * weights

这个例子中没有显式的拆分拼接，但view操作本质上也是一种维度重组。如果要用拆分拼接来实现，可以写成：

# 等价但低效的写法 parts = torch.split(x, 1, dim=1) weighted = [p * w for p, w in zip(parts, weights.split(1, dim=1))] result = torch.cat(weighted, dim=1)

显然，直接用广播乘法更高效。这个对比说明：拆分拼接虽然灵活，但不是所有维度操作都需要用它，很多时候广播和view是更好的选择。

6.4 自定义数据集中的变长序列处理

在NLP任务中，一个batch内的序列长度可能不同。常见的做法是用pad_sequence补齐，然后用pack_padded_sequence处理。但如果不使用这些高级API，也可以手动拆分拼接：

def collate_fn(batch): """手动实现变长序列的拼接""" sequences = [item['input_ids'] for item in batch] lengths = [len(s) for s in sequences] max_len = max(lengths) # 补齐 padded = [] for s in sequences: pad_len = max_len - len(s) if pad_len > 0: s = torch.cat([s, torch.zeros(pad_len, dtype=s.dtype)]) padded.append(s) # 堆叠成batch return torch.stack(padded, dim=0), torch.tensor(lengths)

这个模板展示了cat（补齐）和stack（组batch）的典型配合使用。

7. 一些零散但实用的经验

关于torch.split和torch.chunk的选择，我再补充一个实际经验：在写需要导出ONNX的模型时，尽量用split而不是chunk。因为chunk的不均匀拆分行为在ONNX导出时可能不被支持，或者导出后的行为与PyTorch不一致。split传入明确的大小列表，导出结果更可预测。

关于torch.cat的性能，有一个容易被忽略的点：当拼接的张量在内存中已经是连续且相邻的时候，PyTorch可能会做优化。但这种情况很少见，大多数时候还是老老实实拷贝。所以不要指望cat是零开销的。

关于torch.stack，有一个使用技巧：当你需要把一批numpy数组转成张量时，先用np.stack再torch.from_numpy，比先逐个转张量再torch.stack要快。因为numpy的stack是在C层面完成的，而PyTorch的stack会逐个处理。

import numpy as np # 较快 arrays = [np.random.randn(3, 4) for _ in range(100)] tensor1 = torch.from_numpy(np.stack(arrays)) # 较慢 tensors = [torch.from_numpy(a) for a in arrays] tensor2 = torch.stack(tensors)

最后分享一个调试技巧：在怀疑拆分拼接出问题时，用torch.equal来验证两种实现是否等价。比如你优化了一段代码，把chunk+cat改成了view+transpose，可以用torch.equal确认结果完全一致。这比肉眼对比形状要可靠得多。

# 验证两种实现是否等价 x = torch.randn(2, 64, 8, 8) result1 = torch.cat(torch.chunk(x, 4, dim=1), dim=1) result2 = x.view(2, 4, 16, 8, 8).transpose(1, 2).reshape(2, 64, 8, 8) print(torch.equal(result1, result2)) # 需要根据实际变换调整

这些经验都是我在实际项目中一点点积累的，有些是踩了坑才记住的，有些是看源码和文档时注意到的。拆分拼接看起来简单，但真正在复杂模型中用好，需要对每个函数的行为边界有清晰的认识。希望这些内容能帮你少走一些弯路。