1. 为什么张量拆分与拼接是绕不开的基本功刚接触PyTorch那会儿我总觉得张量操作里最不起眼的就是torch.split、torch.chunk、torch.cat、torch.stack这几个函数。直到有一次在做一个多分支网络的特征融合模块输入特征图需要按通道切成四份分别送进不同的卷积分支处理完再拼回来。我随手写了torch.chunk(x, 4, dim1)结果通道数不能被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]]) # 沿dim0每份大小为2 parts torch.split(x, 2, dim0) print(len(parts)) # 2 print(parts[0].shape) # torch.Size([2, 6]) print(parts[1].shape) # torch.Size([2, 6]) # 沿dim1每份大小为3 parts torch.split(x, 3, dim1) 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) # 沿dim1切成3份 parts torch.chunk(x, 3, dim1) print([p.shape for p in parts]) # [torch.Size([4, 2]), torch.Size([4, 2]), torch.Size([4, 2])] # 沿dim1切成4份6不能被4整除 parts torch.chunk(x, 4, dim1) 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) # 沿dim0拆开 rows torch.unbind(x, dim0) print(len(rows)) # 4 print(rows[0].shape) # torch.Size([6]) # 沿dim1拆开 cols torch.unbind(x, dim1) print(len(cols)) # 6 print(cols[0].shape) # torch.Size([4])unbind的一个典型使用场景是处理序列数据时需要逐时间步操作。但要注意unbind返回的是原张量的视图view对返回的张量做原地修改会影响原张量。这一点和split、chunk一样它们返回的都是视图而非拷贝。2.4 三者的选择依据与常见误用函数语义返回份数不均匀时的行为典型场景split按大小切由大小决定最后一份可能较小精确控制每份大小chunk按份数切尽量接近指定份数前面大后面小快速均分unbind完全拆开等于维度长度无逐元素/逐步处理我个人的选择习惯是需要精确控制每份大小时用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) # 沿dim1拼接 c torch.cat([a, b], dim1) print(c.shape) # torch.Size([2, 7]) # 沿dim0拼接 d torch.cat([a, a], dim0) print(d.shape) # torch.Size([4, 3])cat的核心约束是除了拼接维度外其他所有维度的大小必须完全一致。比如上面沿dim1拼接时dim0的大小必须都是2。如果不一致会直接报错。a torch.randn(2, 3) b torch.randn(3, 4) try: torch.cat([a, b], dim1) 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.这个报错信息其实很清晰它告诉你除了dim1之外其他维度期望的大小和实际得到的大小。排查时重点看报错里说的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) # 沿dim0堆叠 s torch.stack([a, b, c], dim0) print(s.shape) # torch.Size([3, 3, 4]) # 沿dim1堆叠 s torch.stack([a, b, c], dim1) print(s.shape) # torch.Size([3, 3, 4]) # 沿dim2堆叠 s torch.stack([a, b, c], dim2) 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的取值范围是[-(N1), 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, dim1) # (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], dim1) 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, padding1) self.branch3 nn.Conv2d(branch_channels, out_channels // 4, 5, padding2) self.branch4 nn.Conv2d(branch_channels, out_channels // 4, 7, padding3) def forward(self, x): # 沿通道维度分成4份 parts torch.chunk(x, 4, dim1) # 分别处理 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], dim1)这个模式在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, dim0) 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 tensors size and stride... # 正确做法 parts[0].contiguous().view(-1)cat和stack对非连续张量是兼容的它们会正确处理步长。但如果你在拼接后要做view最好在拼接前或拼接后调用.contiguous()避免运行时错误。5. 实战中容易踩的坑与排查思路5.1 维度不匹配报错的快速定位方法维度不匹配是拆分拼接中最常见的错误。PyTorch的报错信息通常会告诉你哪个维度不匹配但有时候信息比较长容易看漏。我的排查习惯是先看报错最后一行通常会给出期望大小和实际大小。打印所有输入张量的shape在拼接前加一行print([t.shape for t in tensors])。检查拼接维度是否正确特别是用了负数索引时。# 一个典型的排查辅助函数 def safe_cat(tensors, dim0): 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, dimdim)这个辅助函数在调试阶段很有用能快速定位是哪个维度出了问题。5.2 拆分后形状不确定导致的后续错误chunk的不均匀拆分是隐蔽性很强的坑。比如你写了一个模块假设输入通道是64用chunk(x, 3, dim1)分成3份你以为每份是21、21、22或者类似但实际上ceil(64/3)22所以是22、22、20。如果你的后续代码假设每份都是22就会出错。x torch.randn(2, 64, 8, 8) parts torch.chunk(x, 3, dim1) print([p.shape[1] for p in parts]) # [22, 22, 20]避免这个坑的最好方法是要么确保能整除要么用split传入明确的列表。# 明确指定每份大小 parts torch.split(x, [22, 22, 20], dim1)5.3 原地操作与视图共享的内存陷阱split、chunk、unbind返回的都是视图和原张量共享内存。如果你修改了拆分结果原张量也会被修改。x torch.arange(12).reshape(3, 4) parts torch.split(x, 1, dim0) 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, dim0)]5.4 在GPU上的拆分拼接性能实测在GPU上拆分和拼接的开销比CPU上更值得关注因为GPU的内存分配和内核启动都有额外开销。我做了一个简单的测试import torch import time device torch.device(cuda) x torch.randn(64, 256, 32, 32, devicedevice) # 测试chunk cat的开销 torch.cuda.synchronize() start time.time() for _ in range(1000): parts torch.chunk(x, 4, dim1) y torch.cat(parts, dim1) torch.cuda.synchronize() print(fchunkcat耗时: {time.time() - start:.4f}s) # 对比直接clone torch.cuda.synchronize() start time.time() for _ in range(1000): y x.clone() torch.cuda.synchronize() print(fclone耗时: {time.time() - start:.4f}s)实测下来chunkcat的耗时大约是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, dim1): 将序列按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, dtypex.dtype, devicex.device) x torch.cat([x, padding], dimdim) # 用split切分 segments torch.split(x, seg_len, dimdim) 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, dimdim)6.3 通道注意力的拆分-加权-拼接通道注意力机制如SE Block中经常需要把特征图拆开、加权、再拼回去class ChannelAttention(nn.Module): def __init__(self, channels, reduction16): 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, dim1) weighted [p * w for p, w in zip(parts, weights.split(1, dim1))] result torch.cat(weighted, dim1)显然直接用广播乘法更高效。这个对比说明拆分拼接虽然灵活但不是所有维度操作都需要用它很多时候广播和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, dtypes.dtype)]) padded.append(s) # 堆叠成batch return torch.stack(padded, dim0), 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来验证两种实现是否等价。比如你优化了一段代码把chunkcat改成了viewtranspose可以用torch.equal确认结果完全一致。这比肉眼对比形状要可靠得多。# 验证两种实现是否等价 x torch.randn(2, 64, 8, 8) result1 torch.cat(torch.chunk(x, 4, dim1), dim1) result2 x.view(2, 4, 16, 8, 8).transpose(1, 2).reshape(2, 64, 8, 8) print(torch.equal(result1, result2)) # 需要根据实际变换调整这些经验都是我在实际项目中一点点积累的有些是踩了坑才记住的有些是看源码和文档时注意到的。拆分拼接看起来简单但真正在复杂模型中用好需要对每个函数的行为边界有清晰的认识。希望这些内容能帮你少走一些弯路。
1. 为什么张量拆分与拼接是绕不开的基本功刚接触PyTorch那会儿我总觉得张量操作里最不起眼的就是torch.split、torch.chunk、torch.cat、torch.stack这几个函数。直到有一次在做一个多分支网络的特征融合模块输入特征图需要按通道切成四份分别送进不同的卷积分支处理完再拼回来。我随手写了torch.chunk(x, 4, dim1)结果通道数不能被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]]) # 沿dim0每份大小为2 parts torch.split(x, 2, dim0) print(len(parts)) # 2 print(parts[0].shape) # torch.Size([2, 6]) print(parts[1].shape) # torch.Size([2, 6]) # 沿dim1每份大小为3 parts torch.split(x, 3, dim1) 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) # 沿dim1切成3份 parts torch.chunk(x, 3, dim1) print([p.shape for p in parts]) # [torch.Size([4, 2]), torch.Size([4, 2]), torch.Size([4, 2])] # 沿dim1切成4份6不能被4整除 parts torch.chunk(x, 4, dim1) 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) # 沿dim0拆开 rows torch.unbind(x, dim0) print(len(rows)) # 4 print(rows[0].shape) # torch.Size([6]) # 沿dim1拆开 cols torch.unbind(x, dim1) print(len(cols)) # 6 print(cols[0].shape) # torch.Size([4])unbind的一个典型使用场景是处理序列数据时需要逐时间步操作。但要注意unbind返回的是原张量的视图view对返回的张量做原地修改会影响原张量。这一点和split、chunk一样它们返回的都是视图而非拷贝。2.4 三者的选择依据与常见误用函数语义返回份数不均匀时的行为典型场景split按大小切由大小决定最后一份可能较小精确控制每份大小chunk按份数切尽量接近指定份数前面大后面小快速均分unbind完全拆开等于维度长度无逐元素/逐步处理我个人的选择习惯是需要精确控制每份大小时用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) # 沿dim1拼接 c torch.cat([a, b], dim1) print(c.shape) # torch.Size([2, 7]) # 沿dim0拼接 d torch.cat([a, a], dim0) print(d.shape) # torch.Size([4, 3])cat的核心约束是除了拼接维度外其他所有维度的大小必须完全一致。比如上面沿dim1拼接时dim0的大小必须都是2。如果不一致会直接报错。a torch.randn(2, 3) b torch.randn(3, 4) try: torch.cat([a, b], dim1) 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.这个报错信息其实很清晰它告诉你除了dim1之外其他维度期望的大小和实际得到的大小。排查时重点看报错里说的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) # 沿dim0堆叠 s torch.stack([a, b, c], dim0) print(s.shape) # torch.Size([3, 3, 4]) # 沿dim1堆叠 s torch.stack([a, b, c], dim1) print(s.shape) # torch.Size([3, 3, 4]) # 沿dim2堆叠 s torch.stack([a, b, c], dim2) 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的取值范围是[-(N1), 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, dim1) # (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], dim1) 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, padding1) self.branch3 nn.Conv2d(branch_channels, out_channels // 4, 5, padding2) self.branch4 nn.Conv2d(branch_channels, out_channels // 4, 7, padding3) def forward(self, x): # 沿通道维度分成4份 parts torch.chunk(x, 4, dim1) # 分别处理 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], dim1)这个模式在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, dim0) 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 tensors size and stride... # 正确做法 parts[0].contiguous().view(-1)cat和stack对非连续张量是兼容的它们会正确处理步长。但如果你在拼接后要做view最好在拼接前或拼接后调用.contiguous()避免运行时错误。5. 实战中容易踩的坑与排查思路5.1 维度不匹配报错的快速定位方法维度不匹配是拆分拼接中最常见的错误。PyTorch的报错信息通常会告诉你哪个维度不匹配但有时候信息比较长容易看漏。我的排查习惯是先看报错最后一行通常会给出期望大小和实际大小。打印所有输入张量的shape在拼接前加一行print([t.shape for t in tensors])。检查拼接维度是否正确特别是用了负数索引时。# 一个典型的排查辅助函数 def safe_cat(tensors, dim0): 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, dimdim)这个辅助函数在调试阶段很有用能快速定位是哪个维度出了问题。5.2 拆分后形状不确定导致的后续错误chunk的不均匀拆分是隐蔽性很强的坑。比如你写了一个模块假设输入通道是64用chunk(x, 3, dim1)分成3份你以为每份是21、21、22或者类似但实际上ceil(64/3)22所以是22、22、20。如果你的后续代码假设每份都是22就会出错。x torch.randn(2, 64, 8, 8) parts torch.chunk(x, 3, dim1) print([p.shape[1] for p in parts]) # [22, 22, 20]避免这个坑的最好方法是要么确保能整除要么用split传入明确的列表。# 明确指定每份大小 parts torch.split(x, [22, 22, 20], dim1)5.3 原地操作与视图共享的内存陷阱split、chunk、unbind返回的都是视图和原张量共享内存。如果你修改了拆分结果原张量也会被修改。x torch.arange(12).reshape(3, 4) parts torch.split(x, 1, dim0) 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, dim0)]5.4 在GPU上的拆分拼接性能实测在GPU上拆分和拼接的开销比CPU上更值得关注因为GPU的内存分配和内核启动都有额外开销。我做了一个简单的测试import torch import time device torch.device(cuda) x torch.randn(64, 256, 32, 32, devicedevice) # 测试chunk cat的开销 torch.cuda.synchronize() start time.time() for _ in range(1000): parts torch.chunk(x, 4, dim1) y torch.cat(parts, dim1) torch.cuda.synchronize() print(fchunkcat耗时: {time.time() - start:.4f}s) # 对比直接clone torch.cuda.synchronize() start time.time() for _ in range(1000): y x.clone() torch.cuda.synchronize() print(fclone耗时: {time.time() - start:.4f}s)实测下来chunkcat的耗时大约是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, dim1): 将序列按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, dtypex.dtype, devicex.device) x torch.cat([x, padding], dimdim) # 用split切分 segments torch.split(x, seg_len, dimdim) 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, dimdim)6.3 通道注意力的拆分-加权-拼接通道注意力机制如SE Block中经常需要把特征图拆开、加权、再拼回去class ChannelAttention(nn.Module): def __init__(self, channels, reduction16): 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, dim1) weighted [p * w for p, w in zip(parts, weights.split(1, dim1))] result torch.cat(weighted, dim1)显然直接用广播乘法更高效。这个对比说明拆分拼接虽然灵活但不是所有维度操作都需要用它很多时候广播和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, dtypes.dtype)]) padded.append(s) # 堆叠成batch return torch.stack(padded, dim0), 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来验证两种实现是否等价。比如你优化了一段代码把chunkcat改成了viewtranspose可以用torch.equal确认结果完全一致。这比肉眼对比形状要可靠得多。# 验证两种实现是否等价 x torch.randn(2, 64, 8, 8) result1 torch.cat(torch.chunk(x, 4, dim1), dim1) result2 x.view(2, 4, 16, 8, 8).transpose(1, 2).reshape(2, 64, 8, 8) print(torch.equal(result1, result2)) # 需要根据实际变换调整这些经验都是我在实际项目中一点点积累的有些是踩了坑才记住的有些是看源码和文档时注意到的。拆分拼接看起来简单但真正在复杂模型中用好需要对每个函数的行为边界有清晰的认识。希望这些内容能帮你少走一些弯路。
