如何用3D张量索引2D张量？PyTorch中的索引与squeeze疑问

一、为什么B[A]会得到形状(1,4,4,2)的张量？ #

• 先明确两个张量的基础信息：
  • B是torch.eye(2)生成的2×2对角张量，即[[1,0],[0,1]]，每个索引i对应的B[i]是长度为2的行向量（比如B[0]是[1,0]，B[1]是[0,1]）。
  • A是形状(1,4,4)的掩码张量，所有元素都是0或1，相当于一个三维的索引集合。
• PyTorch的高级索引规则：当用一个张量作为索引去取另一个张量的元素时，会把索引张量中的每个元素，替换成被索引张量中对应位置的行（或维度匹配的元素）。
• 具体到这个场景：A中每个位置的0或1，都会被替换成B中对应索引的行向量（长度2）。原来A的形状是(1,4,4)，每个元素扩展成长度2的向量后，最终形状就变成了(1,4,4,2)——相当于在原张量的最后新增了一个维度，维度大小等于B的列数（也就是2）。

二、为什么DiceLoss代码中需要用squeeze(1)？ #

• 先看代码逻辑：torch.eye(num_classes)[y_true.squeeze(1)]的目的是把分割标签转换成one-hot编码。
• 实际分割任务中，输入的标签张量y_true通常会带有一个冗余的通道维度，也就是形状为(b,1,h,w)（b是batch size，1是单通道，h、w是特征图尺寸），而非你描述的(b,h,w)——大概率是你忽略了这个初始的通道维度。
• squeeze(1)的作用是去掉这个冗余的通道维度：把(b,1,h,w)转换成(b,h,w)，这样每个位置的标签值（类别索引）才能直接作为索引去取torch.eye(num_classes)的行向量。如果不去掉这个维度，索引张量形状是(b,1,h,w)，最终得到的one-hot张量会多一个不必要的维度(b,1,h,w,c)，后续计算DiceLoss时会和预测张量（形状(b,c,h,w)）的维度不匹配，导致报错。
• 举个例子：假设num_classes=3，y_true原形状是(2,1,4,4)，squeeze(1)后变成(2,4,4)，用torch.eye(3)索引后得到(2,4,4,3)，转置后就能和(2,3,4,4)的预测张量对齐计算损失。

补充：相关原理视频推荐 #

可以搜索PyTorch官方教程中的「高级索引」章节视频，或者国内视频平台上的PyTorch张量索引实战讲解，重点看整数索引和one-hot编码转换部分的内容。

内容的提问来源于stack exchange，提问作者wangxiang