PyTorch张量操作精讲：索引分片、合并与维度调整-北京尧图网络

开头先说一下为什么要把这几个操作单独拿出来写一篇。不管是做CV还是做NLP也不管是在搭网络还是在写数据加载逻辑PyTorch里最绕不开的就是张量操作。我见过不少人背了一遍API就开始写模型结果一遇到形状对不上、维度爆炸、广播规则搞不清就开始懵。这篇笔记就围绕索引分片、合并和维度调整这三块把常用的方法原理、实际用法和你以后一定会踩的坑讲清楚。内容按“先定位元素—再组合数据—最后改形状”的顺序推进新手可以顺着学有基础的朋友直接跳到自己薄弱的部分查漏补缺。1. 索引与分片先掌握张量的“定位”能力1.1 从Python列表说起NumPy风格还是PyTorch风格PyTorch的张量索引整体延续了NumPy的语法风格。定义一个二维张量import torch x torch.arange(20).reshape(4, 5) # tensor([[ 0, 1, 2, 3, 4], # [ 5, 6, 7, 8, 9], # [10, 11, 12, 13, 14], # [15, 16, 17, 18, 19]])最基本的整数索引和Python列表类似x[0]拿到第一行x[0, 1]拿到第一行第二列的元素从0开始计数。但这里有一个非常容易忽略的区别x[0]返回值是张量而不是列表视图它的维度是(5,)。如果你写x[0][1]虽然也能取到元素但写代码时建议把索引放在同一对方括号里x[0, 1]。两种写法结果一样但后者更符合张量的语义切分的表达力也更强。Python列表的切片是list[start:stop:step]PyTorch张量的切片语法同理只是可以针对每个维度各写一段切片。比如x[1:, 2:] # 从第1行到最后一行、第2列到最后一列 x[:, ::2] # 所有行、从第0列开始每隔一列取一列这里要注意一个和Python列表的差异切片结果和原张量共享内存。也就是说对切片结果做原地修改会影响原始数据。这一点既像NumPy又比Python列表更“危险”。很多人写数据增强时在切片上做归一化以为只是改了副本结果原始数据也被改了再拿来算损失就全错了。除了切片后面讲的view、transpose等操作也有这类共享内存问题建议养成一个习惯只要打算后续修改数据就显式调一次.clone()。1.2 花式索引Fancy Indexing的原理与踩坑点花式索引是用张量/列表作为索引去取数据。它可以按行取、按列取也可以同时按行列取。直接看代码indices torch.tensor([0, 2, 3]) x[indices] # 取第0、2、3行 x[:, indices] # 取所有行的第0、2、3列 x[indices, indices] # 取(0,0), (2,2), (3,3)三个位置的元素这里的重点在最后一行x[indices, indices]并不是取“行索引和列索引的笛卡尔积”而是把两个索引张量按位置一一配对。想要笛卡尔积得用组合索引或者torch.meshgrid先构建对应的矩阵。花式索引的结果和切片不一样它会把数据复制一份得到新张量不和原张量共享内存。这点刚入门很容易搞混。切片是“视图”修改会影响原数据花式索引是“拷贝”修改不影响原数据。还有一个小细节花式索引传Python列表和传张量的差别。x[[0, 2, 3]]和x[torch.tensor([0, 2, 3])]拿到的数据一样但结果的类型和梯度传播路径是相同的。真正要注意的是当索引值超出了张量维度范围运行时会直接报错IndexError: index 5 is out of bounds这一点比NumPy更严格NumPy在部分版本下会产生警告但不一定报错。PyTorch的严格模式其实是好事越早发现问题越不容易留下隐藏bug。1.3 布尔索引实战中最好用的筛选工具布尔索引在数据预处理阶段用得非常多。它的核心思路是用条件和原张量形状相同的布尔张量返回所有为True的位置上的元素。几个典型场景x torch.tensor([1.0, -2.0, 3.5, -0.5]) mask x 0 x[mask] # 取所有正数 x[x 0] 0 # 把所有负数原地置为0 # 多维场景找出所有大于10的元素 x2 torch.randint(0, 20, (4, 5)) rows torch.nonzero(x2 10) # 返回满足条件的坐标布尔索引的结果也总是进行数据拷贝。还有一个高频技巧是结合多个条件x[(x 0) (x 10)]注意括号不能省。很多人写成x[x 0 and x 10]直接报错因为Python的and不能直接用在张量上必须逐元素逻辑运算、|。对了想定位坐标时不要用torch.where去“取元素”它更适合做三目运算。要拿到满足条件位置的坐标直接用torch.nonzero返回形状为[非零个数, 维度数]。如果想要“第0维为行索引第1维为列索引”的老式返回结果就调torch.nonzero(..., as_tupleFalse)这个参数默认状态在新版PyTorch里已经调整过使用前看一眼版本说明。2. 合并把多个张量拼在一起2.1 torch.cat拼接沿已有维度把数据接起来合并操作最常用的就是torch.cat和torch.stack两个函数先说说torch.cat。它的作用是在已有的某个维度上把多个张量拼接在一起好比把几段绳子头尾相连。a torch.randn(2, 3) b torch.randn(4, 3) c torch.cat([a, b], dim0) # 形状(6, 3) # 如果做dim1需要两个张量第0维相同 d torch.cat([a, a], dim1) # 形状(2, 6)对于dim0的拼接除了第0维以外的所有维度都必须完全一致对于dim1的拼接除了第1维以外的所有维度都必须一致。这个限制很好理解拼接就像盖楼时把结构件对接只有接口截面一致才能严丝合缝。如果a是(2, 3)b是(2, 4)你硬要做dim1运行时会报错RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 3 but got size 4这段报错信息几乎是所有新手都会遇到的。解决办法就是先检查形状再决定拼接的维度。多数情况下你真正需要的是“把不同样本拼成一个batch”也就是在dim0上拼这时只要保持每个样本的内部维度一致即可。在实际项目中torch.cat最常见的用法是在DataLoader的collate_fn里处理变长数据。比如NLP里一句话七个词另一句话九个词没法直接拼成(2, seq_len, hidden)那就先pad到相同长度再cat。还有一个常见场景是把多尺度特征图在通道维上拼起来U-Net的skip connection就是典型这时就要用dim1。2.2 torch.stack合并凭空多出一个维度torch.stack和torch.cat最大的区别在于stack不会沿着已有维度拼接而是把所有输入张量堆叠在一个新维度上。a torch.zeros(2, 3) b torch.ones(2, 3) c torch.stack([a, b], dim0) # 形状(2, 2, 3) d torch.stack([a, b], dim1) # 形状(2, 2, 3) e torch.stack([a, b], dim2) # 形状(2, 3, 2)以dim0为例新的第0维的尺寸就是输入张量的个数。既然是新维度那么要求所有输入张量的形状完全相同。cat允许某个维度不同stack则要求全部一致。什么时候用stack呢最典型的就是把一组独立维度相同的张量合成batch。比如你的数据是[特征向量]每条形状都是(10,)有32条数据你希望的batch形状是(32, 10)用torch.stack(data_list)直接搞定。如果用torch.cat两个(10,)在dim0上只能得到(20,)维度就少了一层。还有一个小技巧想给张量批量增加一个维度例如把(3, 4, 5)变成(3, 1, 4, 5)时也可以用torch.stack把同一个张量放两次再选其中一个切片即可但更干净的做法是直接unsqueeze后面专门讲。2.3 合并操作的常见报错与处理思路合并时比较常见的两个报错第一个是“形状不一致”报错。cat要求非目标维度完全一致stack要求所有维度完全一致。处理思路就一句话看报错里提到的维度数字再回去打印两个张量的.shape逐一比对。第二个是“数据类型不一致”报错。合并时PyTorch要求参与操作的所有张量dtype必须一样。比如一个张量是torch.float32另一个是torch.float64直接报错。这时要么统一转成float32要么都转成float64在CPU上还会涉及内存翻倍的问题尽量统一用float32。另外还要注意一个容易忽略的坑合并操作不是原地修改不会改变原张量。很多人写torch.cat([a, b], dim0)然后打印a.shape发现没变化以为拼接失败了。其实cat返回新张量必须用c torch.cat([a, b], dim0)这种形式接住返回值。这个概念在PyTorch里无处不在很多操作都不会原地改输入不仅限于合并。3. 维度调整改变张量的形状3.1 view与reshape的关系与本质区别维度调整是PyTorch里最让人迷惑的部分尤其是view和reshape的区别。简单说view要求张量在内存中满足“连续”contiguous条件reshape不要求reshape在底层会自动复制数据使其连续。先看代码x torch.arange(12) y x.view(3, 4) # 形状(12,) - (3, 4) z x.reshape(3, 4) # 同样结果但处理机制不同如果张量已经是连续的两者没有任何区别。但当张量做过transpose、permute等操作后内存布局会变得不连续此时调用view会直接报错RuntimeError: view size is not compatible with input tensors size and stride (...)遇到这种报错最简单的处理办法是x.contiguous().view(...)或者直接用reshape。这里的关键是理解“逻辑形状”和“内存布局”的区别。你可以把张量想象成一本页码连续的字典view只是改变了你读字典的方式但要求页码必须仍然从0到N连续排列一旦你调了transpose把书的章节顺序打乱了页码对不上view就失效了必须先把书重新整理成新的连续页码contiguous再操作。我在项目里一直保持这样的习惯只要不是对性能极其敏感的核心循环优先用reshape而不是view。因为reshape的语义更接近“我要一个形状上的新张量”不用操心底层连续性问题。但如果你做张量并行或者高频处理就要注意reshape在非连续场景下会引入一次拷贝这个开销在某些场景下不可小视。3.2 permute与transpose维度怎么换才不丢数据transpose和permute都是交换维度但transpose一次只交换两个维度permute可以任意排列所有维度。看示例x torch.randn(2, 3, 4) x.transpose(0, 2) # 交换0维和2维形状变成(4, 3, 2) x.permute(2, 1, 0) # 把原2维放前面、原1维放中间、原0维放最后形状同样(4, 3, 2)如果只是交换两个维度transpose(0, 2)和permute(2, 1, 0)结果不完全相同permute实际上可以同时完成多维交换最终形状取决于你给出的轴顺序。记住一个口诀permute(原维度的新顺序)传入的参数表示“新张量的第i维对应原张量的第几个维”。从数据角度看permute和transpose都不会改变数据的值只是改变了读取顺序。它们返回的往往是不连续的张量所以在后续接view前记得先contiguous()。transpose还有一个常见误用x.T。T在PyTorch里是2D张量的转置快捷写法在更高维张量上x.T只能做维度完全反转效果等同于x.permute(*reversed(range(x.dim())))。如果你只想交换两个维度用transpose更直观。在图像处理里transpose是最常见的“通道转换”工具。OpenCV读出来是H, W, C128, 128, 3要转成PyTorch模型需要的C, H, W3, 128, 128一行img.transpose(2, 0, 1)搞定。3.3 unsqueeze、squeeze与flatten的使用场景unsqueeze是在指定位置插入尺寸为1的新维度squeeze是删除所有尺寸为1的维度或指定删除某个维度。flatten则是把连续若干维压成一维。x torch.randn(3, 4) x.unsqueeze(0) # 形状(1, 3, 4)类似增加batch维度 x.unsqueeze(1) # 形状(3, 1, 4) x.unsqueeze(2) # 形状(3, 4, 1) x.squeeze(0) # 如果第0维是1才会删除否则不变 x.reshape(-1) # 展平成一维等价于flattenunsqueeze在深度学习里有个极为经典的场景单条样本转batch。预测时你有一张图片模型要求输入(N, C, H, W)但当前只有(C, H, W)直接img.unsqueeze(0)就能得到(1, C, H, W)。flatten和reshape(-1)的结果通常是一样的但flatten可以指定起止维度。nn.Flatten(start_dim1)就是把从第1维开始往后全部压成一维比如(32, 3, 224, 224)变成(32, 3 * 224 * 224)。全连接层前接特征图时nn.Flatten()是标配。还有一个细节squeeze不加参数时会删除所有尺寸为1的维度。这可能会引入一个隐藏bug——如果你本意只是删掉某一个维度但其他维度恰好也是1结果就会比预期多删。所以需要精确控制时一定要传维度参数x.squeeze(0)。4. 综合实战维度变换在真实项目中的套路4.1 典型流程从加载数据到送入模型把上面三个部分串起来看一个完整的例子。假设你在做一个人脸关键点检测任务数据加载器返回的样本形状是这样的图像(H128, W128, C3)OpenCV读入通道在最后关键点(68, 2)68个点的x/y坐标模型输入要求是图像(batch, C3, H128, W128)关键点标签要求是(batch, 68*2)。import cv2 import torch # 模拟一张图像和关键点 img cv2.imread(face.jpg) # (128, 128, 3) pts torch.randn(68, 2) # (68, 2) # 通道调整HWC - CHW img_t torch.from_numpy(img).permute(2, 0, 1) # (3, 128, 128) # 单样本转batch img_batch img_t.unsqueeze(0) # (1, 3, 128, 128) # 关键点展平 pts_batch pts.reshape(1, -1) # (1, 136) # 多个样本合并成一个大batch用torch.cat batch_imgs torch.cat([img_batch, img_batch], dim0) # (2, 3, 128, 128)这个流程里你同时用到了permute换轴、unsqueeze加维度、reshape展平和cat合并。写模型时会发现几乎每个数据加载流程都会重复这几个操作。把这些套路记住比死背API有用得多。4.2 从PyTorch到ONNX维度兼容问题很多模型训练完之后要导出ONNX热词里也有大量pytorch转onnx的搜索这时维度调整就显得更重要了。ONNX导出时会固定输入张量的shape但允许动态维度。如果你在模型里使用了不规范的view、reshape导出的ONNX计算图里生成的Transpose、Reshape节点就可能会多出很多推理框架如OnnxRuntime跑起来性能就差。实践上建议模型内部尽量用reshape代替view但在知道张量连续时也用view保证效率动态维度场景比如可变batch size用torch.onnx.export时设置动态轴如果模型里包含flatten导出的图中通常会有对应的Reshape节点这是正常的但要注意某些推理硬件对动态reshape支持不好能固定shape就固定。我踩过的一个实际坑是模型里用了x.view(x.size(0), -1)当batch size变化时ONNX导出没问题但在一些NPU设备上推理会报错。换成x.reshape(x.size(0), -1)后同样报错最后改成显式x.contiguous().view(...)才稳定。不是view本身有问题而是部分推理框架对“非连续视图”支持有限导出前最好在模型里显式保证连续性。4.3 广播机制的联动维度调整和加法一起用张量做加减法时维度不同但满足广播规则也能操作。广播规则概括起来是从最后一个维度依次往前比对两个维度相同或者其中一个是1就能对齐。例如(3, 1)和(1, 4)可以相加得到(3, 4)。这个规则和维度调整经常配合使用。比如你想给一个batch的每个样本加一个bias向量data torch.randn(32, 10) # batch, feature bias torch.randn(10) # 一维向量 out data bias # 自动广播(32, 10) (10,) - (32, 10)但如果bias的维度是(10, 1)就不能直接加了需要bias.squeeze(1)或bias.reshape(10)。反过来如果特征是二维特征图(32, 3, 224, 224)要加一个针对通道的均值(3,)就要把它变成(1, 3, 1, 1)再相加常用写法是mean[None, :, None, None]等价于mean.unsqueeze(0).unsqueeze(2).unsqueeze(3)。广播虽然方便但也是最容易出隐藏bug的地方。特别是两个张量形状不匹配但广播后得到意外结果时不会报错只会输出一个更大形状的张量后续计算看着没问题实际已经错了。遇到这种情况的最快排查方式在、*操作后立刻打印out.shape只花一行代码的时间能省一整晚的定位时间。5. 常见报错与排查心得5.1 报错信息速查表报错信息关键词原因处理方式Sizes of tensors must matchcat拼接维度不一致打印两个shape确认非拼接维度完全一致view size is not compatible对非连续张量用了view先contiguous()再view或直接用reshapeindex 5 is out of bounds for dimension索引超出范围检查索引张量最大值与对应维度大小Expected all tensors to be on the same device张量跨设备统一to(device)Expected scalar type Float but found Doubledtype不匹配统一.float()或.double()Boolean indexing requires matching shapes布尔索引mask尺寸不匹配检查mask与目标张量形状这张表基本覆盖了索引、合并、维度调整这三类操作的大多数报错。记牢关键词之后能少走很多弯路。5.2 自动求导背景下的维度陷阱在requires_gradTrue的张量上做维度调整大多数情况下梯度能正常传播但有几个坑要注意。第一不是所有操作都能保持梯度。torch.nonzero和布尔索引这类操作虽然可以用于前向计算但涉及到的索引是不可微的。如果它出现在网络中间梯度回传会在这一步断掉。例如x torch.randn(10, requires_gradTrue) mask x 0 y x[mask].sum() y.backward() # x.grad不会是所有位置都有值这里x.grad只在mask为True的位置上有梯度其他位置为None。如果后续代码对这个梯度做别的运算很容易出问题。第二inplace操作要极度小心。x[0] 1这种原地修改在requires_gradTrue时可能导致梯度计算错误甚至直接报错a leaf Variable that requires grad is being used in an in-place operation。建议不要对叶子张量做任何原地修改。第三view和reshape在梯度回传上通常没有问题因为它们本质上是同一个张量的不同视图梯度会正确地映射回去。但如果你在一个view的结果上做了原地修改再对原张量做计算梯度方向就会乱套。这种情况最好的防御手段是不修改、不多想真要修改先clone()。5.3 内存拷贝与性能注意事项维度调整带来的隐秘性能问题主要体现在内存拷贝上。transpose、permute返回非连续张量contiguous()会触发拷贝flatten在非连续张量上也可能拷贝reshape在非连续时同样拷贝。写高性能推理代码时尽量让数据从头到尾保持连续。一个实用建议数据预处理阶段尽量一次性调整好布局比如直接转成C, H, W不要在模型里反复permute再contiguous。PyTorch模型内部可以使用ModuleList加自定义的Sequential层来固定布局这样对整个pipeline的可读性和性能都有帮助。另一个常见困惑是clone、detach、copy_的区别。clone会复制张量并保留梯度计算图detach会断开梯度连接但共享内存copy_是原地拷贝数据。在维度调整场景中如果你只要一个不含梯度关系的新张量最安全的方式是new_x x.detach().clone()如果不需要梯度也可以用x.detach()直接获得new view再决定要不要clone。这个组合拳在写数据增强、做可视化、存特征时都非常实用。5.4 写一个minimal案例来验证维度操作排查维度问题时我的习惯是从小规模数据开始验证不直接在大张量上试。比如写这样一个最小案例def debug_dim(fn): a torch.randn(2, 3, 4) b torch.randn(2, 3, 4) try: out fn(a, b) print(OK, shape:, out.shape) except Exception as e: print(Error:, e) debug_dim(lambda a, b: torch.cat([a, b], dim1)) # 结果(2, 6, 4)这样能快速验证某个操作的维度语义也方便在Stack Overflow或GitHub Issue里描述问题时贴出来别人一看就懂。毕竟维度调整这种问题代码比文字解释快得多。还有一个小技巧用x.stride()查看张量的内存步幅。很多维度错误通过stride一眼就能看出来比如做了transpose之后stride会变成反向的。看到stride和我们预期不一致就知道当前张量是不连续的后面该不该调contiguous心里就有数了。这个函数在常规教程里提得不多但排查维度问题非常有用。最后再分享一个心得。我在项目里常用的做法是给每个关键张量写一个形状注释类似# x: (batch, time, feature) - (batch, feature, time) x x.transpose(1, 2)虽然看起来多打几个字但调试时能省很多时间。尤其当你拿到别人写的模型代码满屏都是permute和reshape没有注释就只能靠猜。形状就是张量的“身份证”写清楚它比什么都强。

PyTorch张量操作精讲：索引分片、合并与维度调整

相关推荐

5MW永磁直驱风电1200V直流并网Simulink仿真模型搭建与调试

MiMo-V2.6：无奖励函数的自改进强化学习架构

Agent-Reach：轻量级CLI大模型路由调度器

更多资讯

领域特定评估实战：用 Argilla、Distilabel 与 LightEval 构建考试问答评估流水线（smol-course）

Apache Storm 集群安全加固实战：从 OS 层防护到 Kerberos 认证与 ACL 授权

CMake FindOpenCL 模块全解析：从 find_package 到 OpenCL::OpenCL 导入目标

YOLO船舶检测实战：数据集解析与训练避坑指南

题解：洛谷 P14361 [CSP-S 2025] 社团招新

U-Boot Kbuild深度解析：从零构建RV1106移植的四大核心步骤

今日更新

PHP htmlentities()函数用法讲解

AI Agent工程实战：从七要素到七个决策点的系统设计指南

AI入局芯片设计验证：验证收敛50倍加速的真相与应对

本周精选

CMOS设计三大物理效应：体效应、沟道长度调制与亚阈值导电

微信小程序在线阅读毕设全解析：Java后端+MySQL+避坑实战

Android音频问题五层分析法：从App到Codec的系统化诊断

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？