★ FEATURED ARTICLE 1. 为什么这个转换会让人绕进去先说个我自己的经历。有段时间我在调一个文本分类模型输出层后面接了一个形状为(batch_size, seq_len, num_labels)的张量训练标签用的是 one-hot 编码。到了求准确率的环节我下意识写了个torch.argmax(logits, dim1)结果跑出来准确率惨不忍睹半天查不到原因。后来把张量形状打印出来才反应过来这个场景里类别轴根本不在 dim1而是在 dim2。这类问题几乎每隔一段时间就会在社区里看到有人问一遍torch.argmax的dim1到底取的哪个轴为什么 one-hot 编码转整数标签时都是argmax(dim1)这两个问题的本质其实是同一个——one-hot 的类别维度究竟摆在张量的哪个位置以及 argmax 在该维度上压缩时到底在做什么。这篇文章就是把我这些年踩过的、看别人踩过的相关坑系统梳理一遍。读完你不会再纠结dim1还是dim2而且能顺手搞懂argmax与 one-hot、整数标签之间那条“编码—解码”的对应关系。内容面向刚接触 PyTorch 的入门者也适合写过一段时间但偶尔被维度绕晕的进阶玩家——如果你曾经在argmax上调试超过半小时那这篇大概率能帮你省下下一次的半小时。2. dim 参数不是“选第几行”而是“在哪条轴上做压缩”很多人第一次看到torch.argmax(x, dim1)会产生一个直觉在第二个维度上挑出最大值。这个直觉方向没错但执行细节容易被忽略。我们得先把“张量的轴”这件事说清楚否则后面所有讨论都飘着。2.1 用一个三维张量彻底理解轴假设有一段代码import torch x torch.arange(24).reshape(2, 3, 4) print(x)输出会是一个形状为(2, 3, 4)的张量。第一维的大小是 2第二维是 3第三维是 4。用坐标去理解x[1, 2, 3]这个元素的位置对应“第一维下标 1、第二维下标 2、第三维下标 3”。dim0表示沿着第一维移动、去比较“第一维不同下标、其余维度相同下标”的所有元素。举例x[:, 0, 0]这一串是[0, 12]如果对这整个三维张量做torch.argmax(x, dim0)就在所有“其余两个维度固定、第一维变化”的序列里找最大值下标最后得到的张量形状是(3, 4)——写代码验证一下就是torch.argmax(x, dim0) # 形状(3, 4)因为第一维被压缩没了。同理torch.argmax(x, dim1)返回形状(2, 4)dim2返回(2, 3)。这里有个关键认知argmax 的返回值不是最大值本身而是最大值所在的索引。你在 dim1 上做 argmax得到的是“第二维的哪个下标位置取值最大”。这一步正是 one-hot 转整数标签的核心——one-hot 向量里值为 1 的位置就是类别索引argmax 把这个位置“挖”出来。2.2 二维张量里 dim1 的直觉最常见的分类输出形状是(batch_size, num_classes)。如果 batch_size 等于 4、num_classes 等于 5那么张量的每一行就是一个样本的类别概率分布或者 logits 向量。此时torch.argmax(logits, dim1)等价于问对每一个样本第一维固定在 5 个类别得分里谁最大。返回结果是一个长度为 4 的一维张量每个元素取值范围是 0 到 4。这张二维场景其实是新手最熟悉的。很多人卡住是因为碰到三维、四维张量时脑子里还在沿用“行和列”的直觉结果把 dim 选错。所以下面我会从二维场景开始逐步往后推。3. one-hot 编码与整数标签一对互为逆运算的关系one-hot 编码的本质是用一个长度为num_classes的向量表示一个样本的类别向量中只有一个位置是 1其余全是 0。假设类别有 5 种类别 2 的 one-hot 向量就是[0, 0, 1, 0, 0]。而整数标签则直接用一个非负整数表示比如上面那个就是2。两者之间转换自然有两种方向整数标签 → one-hot构造全零向量在整数标签对应位置置 1。one-hot → 整数标签找到值为 1 的位置下标。torch.argmax(dim1)做的恰好就是第二件事。因为 one-hot 向量里只有一个 1、其他全是 0所以“最大值出现的位置”就是“值 1 的位置”也就是原始整数标签。3.1 为什么偏偏是 dim1这里涉及一个约定俗成的张量布局PyTorch 分类任务里模型的输出通常是(batch_size, num_classes)。沿着 batch 方向移动每次固定一个样本在类别维上取 argmax得到的就是该样本预测出的类别。因此 dim1 对应的是类别轴。注意这个“dim1 是类别轴”并不是数学定理而是一种张量形态约定。可以类比成我们自己定义了一个数据格式第 0 维放样本编号第 1 维放类别编号。只要你在构建张量时遵守这个格式那么argmax(dim1)就永远是对的。一旦有人不小心把张量建成(num_classes, batch_size)那类别轴就变成 dim0 了你还用 dim1 去算结果就会错得非常离谱。验证一下import torch # one-hot 标签形状 (4, 5) one_hot torch.tensor([ [0, 1, 0, 0, 0], [0, 0, 0, 1, 0], [1, 0, 0, 0, 0], [0, 0, 0, 0, 1], ]) labels torch.argmax(one_hot, dim1) print(labels) # tensor([1, 3, 0, 4])结果一眼就能看懂每一行的 1 出现在哪一列输出就是哪一列的下标。3.2 转置陷阱如果 one-hot 按列存放呢有一种情况是数据预处理阶段有人用sklearn.preprocessing.OneHotEncoder默认输出稀疏矩阵再转成数组后形状可能是(num_samples, num_classes)这没问题。但如果你自己手动堆叠向量时按列拼接比如one_hot_T torch.tensor([ [0, 0, 1, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 0, 0], ])这个张量形状是(5, 4)每一列是一个样本的 one-hot。此时取argmax(dim0)才是“沿着类别轴找 1 的位置”得到的结果才是每个样本的整数标签。用argmax(dim1)的话你会得到 4 个结果因为压缩了第二维形状变成(5,)而且含义完全变了。从这个例子能得出一个很实用的自查方法在做任何 argmax 转换之前先打印one_hot.shape确认“类别维”在第几维再去选 dim。永远不要凭“大家都这么写”来盲目决定 dim 的值尤其当你的数据来自不同预处理流程时。4. 从二维到高维分类输出不只有一种形状很多刚接触 NLP 或者多标签任务的人会突然发现模型输出不再是规整的(batch_size, num_classes)而是带上了序列长度、时间步、token 数量等额外维度。这时候dim1还是不是类别轴就得重新审视了。4.1 序列标注场景形状是(batch, seq_len, num_classes)文本分类里的常见做法是直接对整句话做池化或者取[CLS]的向量所以输出还能保持二维。但序列标注、命名实体识别这类任务模型会对每个 token 都输出一个类别分布整体形状是(batch_size, seq_len, num_labels)。此时正确的转换写法是preds torch.argmax(logits, dim2)因为 dim2 才是类别轴。如果你 copy 了一段文本分类的代码写的是dim1那得到的结果就是在“序列长度”这个维度上取最大下标——相当于每个样本只挑出了一个 token 位置而且这个位置的含义完全不是类别预测结果自然全乱。这是我实际发生过的事。那次我把argmax(dim1)的结果直接拿去算 F1打印出来全是超过标签范围的数字排查了很久才发现轴选错了。后来我养成了一个习惯在关键张量运算前加一行注释标注清楚每个维度的语义。# logits 形状: (batch_size8, seq_len128, num_labels12) preds torch.argmax(logits, dim-1)4.2 用 dim-1 替代固定数字这里我给一个很多资深 PyTorch 用户都在用的习惯如果类别轴是最后一个维度直接写dim-1不写dim1、dim2。为什么因为从后往前数更不容易出错而且代码里面所有“最后一维是类别”的模型都可以共用同一个写法。比如一个文本分类模型它的输出在经过某些变化后可能变成(batch_size, num_labels)也可能因为加了 CRF 或额外头变成(batch_size, seq_len, num_labels)。如果用dim-1就不需要关心前面到底有几维只要确定类别在最末尾就能稳定地取到正确结果。# 二维情况 logits_2d torch.randn(4, 5) preds_2d torch.argmax(logits_2d, dim-1) # 等价于 dim1 # 三维情况 logits_3d torch.randn(4, 6, 5) preds_3d torch.argmax(logits_3d, dim-1) # 等价于 dim2这段代码两种形状都能跑且含义都是“对每个样本/token在 5 个类别中取最大者”。4.3 多维 one-hot不需要保持 one-hot 形态有一种情况是模型输出或者标签张量的形状不是(batch, num_classes)而是四个维度例如图像分割场景里标签经常是(batch, height, width)的整数张量模型输出的 logits 是(batch, num_classes, height, width)。这时类别维在 dim1。我要特别提醒这个场景非常容易让人犯迷糊因为一眼看上去“batch 后面跟的确实是 num_classes”而且argmax(dim1)在形式上又恢复了“标准答案”。只不过这里的输出结果是一个(batch, height, width)的整数标签图而不是一个向量。# logits 形状: (batch_size2, num_classes21, height32, width32) seg_preds torch.argmax(logits, dim1) # 返回 (2, 32, 32)这种“类别维排在第 1 维”的布局在视觉任务里很常见。所以再次强调dim 的值不是由“第几个维度”决定的而是由你的数据布局决定的。拿到张量后第一反应应该是看 shape而不是背公式。5. argmax 之外的替代写法以及为什么它们没那么好用one-hot 转整数标签不止argmax一种途径。实际开发中我还见过一些人用torch.where、torch.nonzero、torch.max等方式。这里做一个简单对比方便你根据场景选择。方法代码示例返回内容适用场景argmaxtorch.argmax(x, dim1)索引张量绝大多数分类任务maxtorch.max(x, dim1)最大值和索引namedtuple需要同时拿到分数与索引nonzerotorch.nonzero(x 1)所有非零位置坐标稀疏 one-hot、多标签场景wheretorch.where(x 1)每个非零位置坐标元组调试小批量时偶尔用torch.max的写法长这样values, indices torch.max(one_hot, dim1)它的indices和argmax返回结果完全一致但顺带把最大值也拿回来了。在有些需要同时知道“模型对这个样本的置信度”的场景用torch.max可以少跑一次前向或者少一次索引取值。不过有一个真正的坑torch.max(x, dim1)返回的是一个torch.return_types.max对象如果你直接print它会看到类似torch.return_types.max(valuestensor([...]), indicestensor([...]))的东西。新手容易把这个对象当成普通张量直接用结果报错。稳妥做法是像上面那样解包_, preds torch.max(logits, dim1)至于nonzero更多用在多标签分类里——一个样本可能同时属于多个类别one-hot 不再只有一个 1。对于多标签 one-hot严格说是 multi-hotargmax只能返回其中一个类别而nonzero能把所有类别都列出来multi_hot torch.tensor([[0, 1, 1, 0], [1, 0, 0, 1]]) indices torch.nonzero(multi_hot) print(indices) # tensor([[0, 1], [0, 2], [1, 0], [1, 3]])每个坐标的第一列是样本编号第二列是类别编号。如果想按样本分组取出类别列表可以用labels [row[1].tolist() for row in torch.nonzero(multi_hot).tolist()]不过这样效率一般小批量调试可以跑大数据集还是建议直接用argmax配合掩码处理。6. argmax 与 softmax 的配合别在概率上白做一次指数运算另一个高频相关话题是模型最后一层到底该不该接 softmax再接 argmax很多人会写probs torch.softmax(logits, dim1) preds torch.argmax(probs, dim1)这个写法在结果上没错但属于“白做了一次指数运算”。因为 softmax 是一个严格单调递增的变换也就是说logits 里哪个位置最大softmax 之后哪个位置仍然最大。argmax根本不需要知道具体的概率值它只需要相对大小关系。所以正确的做法是直接在 logits 上做 argmaxpreds torch.argmax(logits, dim1)省掉一次 softmax在 batch 很大的时候能省下一部分计算开销。我记得早期自己写评估循环时每次 forward 出来都习惯性先 softmax 再 argmax后来看 profiling 才发现这里白白多了一整层指数与归一化计算。虽然单次量不大但迭代成千上万次后差距就出来了。6.1 需要 softmax 的场景是什么有一种情况是你需要的不只是预测类别还要置信度分数比如要过滤低置信度的预测结果。此时可以probs torch.softmax(logits, dim-1) max_probs, preds torch.max(probs, dim-1) # 过滤置信度低于 0.8 的预测 mask max_probs 0.8 filtered_preds torch.where(mask, preds, torch.tensor(-1))把低置信度样本的预测标签置为 -1方便后续统一处理。这里同时用到了softmax、max和where算是一个常见的组合拳。注意max的解包顺序第一个返回的是值第二个是索引别搞反。7. 最容易翻车的边界情况与调试技巧平时写代码时除了维度选错还有几个边界状况几乎人人都遇到过我一起列出来。7.1 全零行或全 NaN 张量导致的结果漂移当 one-hot 编码里出现全零行时——比如某条样本的标签确实缺失one-hot 向量全为 0——argmax会返回下标 0。因为最大值 0 出现在第一个位置。这会导致真实类别是缺省预测却变成类别 0而且这个过程完全静默没有报错。这就是一个典型的“不报错但是结果错”的坑。调试的时候可以顺带统计一下row_sums one_hot.sum(dim1) invalid_mask row_sums 0 print(finvalid one-hot rows: {invalid_mask.sum().item()})如果发现存在全零行最好回到数据处理环节决定是丢弃还是填充特殊类别。7.2 同一个最大值出现多次argmax 返回第一次出现的位置假设某一行的 logits 是[1.0, 2.0, 2.0, 0.5]最大值 2.0 同时出现在下标 1 和 2。argmax规定返回第一次出现的下标也就是 1。这个行为在绝大多数场景没有问题但在“并列最大”代表某种平局语义时可能不符合预期。如果业务上需要处理平局你得自己额外判断logits torch.tensor([[1.0, 2.0, 2.0, 0.5]]) max_val, first_idx torch.max(logits, dim1) # 看看有多少位置等于 max_val num_max (logits max_val.unsqueeze(-1)).sum(dim1)当num_max 1时你就可以知道该样本发生了平局再按业务规则处理。7.3 模型输出维度被 squeeze 掉之后的迷茫还有一个高频坑模型 forward 里某个操作不经意间把维度删了。举个例子如果你的模型输出logits原本是(batch_size, 1, num_classes)然后有人在中间加了一行.squeeze(1)它就变成了(batch_size, num_classes)。这本身没问题但如果你把这段代码用在另一个输出(batch_size, seq_len, num_classes)的模型上.squeeze(1)不会删除 seq_len 维度因为那个维度大小不是 1于是后续所有argmax(dim1)全错位。排查这种问题时我的习惯是在 debug 模式下加一段形状断言assert logits.dim() in (2, 3), funexpected logits dim: {logits.shape} if logits.dim() 2: preds torch.argmax(logits, dim-1) elif logits.dim() 3: preds torch.argmax(logits, dim-1)注意这里两种情况其实可以用同一个dim-1搞定但有了断言至少后续人维护代码时能清楚知道模型输出预期是什么形状。7.4 维度含义速查表为了日常翻查方便我把常见张量布局和 argmax 的正确 dim 整理成一个表任务类型常见输出形状类别轴推荐 argmax 写法单标签分类(batch, num_classes)第 1 维dim1或dim-1序列标注(batch, seq_len, num_classes)第 2 维dim2或dim-1图像分割(batch, num_classes, h, w)第 1 维dim1多标签分类(batch, num_classes)第 1 维dim1但不能只用 argmax批次解码(batch, seq_len, vocab_size)第 2 维dim2或dim-1这表不是让你背的而是让你在拿不准的时候心里有个底。核心仍然是先看 shape再确定类别轴的位置最后选择 dim。8. 一次完整的踩坑复盘dim1 如何毁掉一个 NER 实验最后用我开头提到的 NER 实验做一次完整复盘把排查思路串起来。当时情况是这样模型输出 logits 形状是(4, 128, 12)4 是 batch_size128 是 seq_len12 是标签数。我在评估脚本里复制了一段文本分类的代码写的是torch.argmax(logits, dim1)。输出的preds形状是(4, 12)每个元素是 0 到 127 之间的整数看起来像某种 index但和标签完全不匹配。F1 分数跌到 0.1 以下打印几条预测结果后发现数值范围明显异常出现 126、127 这种而真实标签总共才 12 类。检查输入输出 shape 后定位到 argmax 的 dim 选错。修复其实就一行preds torch.argmax(logits, dim-1)但为什么当时没有第一时间发现因为我被“dim1 是类别维”这个惯性带跑了看到 logits 就直接套用完全没核对形状。后来我把评估脚本里所有 argmax 调用都改成了dim-1并要求自己写清楚 shape 注释。从那以后这种维度问题基本没再犯过。复盘得到一条最实用的经验报错不是你唯一的调试信号结果不合理也是一种报错。如果你的预测 index 超出了标签定义范围或者准确率突然崩了先别怀疑模型先去打印张量形状再去看维度轴上每个元素的语义。另外一个小技巧val 阶段可以抽样打印一批“预测标签 输入文本/图像”人眼扫一遍往往比任何指标都能更快暴露问题。我当时就是看到预测标签出现 126 这种数字才意识到不是模型没学好而是轴取错了。9. 写代码时的几条习惯性建议说了这么多最后沉淀几条实操习惯。这些不是理论是我自己在多个项目里试错后留下来的肌肉记忆。第一所有关键张量运算前面写一行形状注释。这行注释不花时间但在三天后回读代码时能救你命。因为你对当时的数据布局记忆会很快消退而对计算机来说形状就是一切。第二能不写固定数字 dim 就不写。如果类别轴总是在最后一维尽量用dim-1。它不是银弹但能大幅度减少维度错位的问题。只有当类别轴明确不在最后一维比如图像分割的输出(batch, num_classes, h, w)才去写dim1。第三one-hot 转整数标签之前先花一秒钟确认 one-hot 是不是真的 one-hot。如果因为某些预处理错误导致某一行有多个 1或者全是 0argmax 的结果会“看起来没问题实际上有问题”。加个 sum 检查不费电。第四把维度语义固化到命名里。比如变量名logits_b_sl_cls虽然丑但比单纯的output或x能传递更多信息。尤其是在多任务、多输出头的模型里这一招能省掉很多沟通成本。回到标题那句话——torch.argmax的dim1与 one-hot 转整数标签的关系说到底就是一件事找到类别轴然后在那条轴上取最大值的下标。类别轴在第几维dim 就是几one-hot 的唯一非零位置就是最大值位置argmax 把它还原成整数标签。这个关系本身不复杂复杂的从来都是张量形状千变万化的现实里如何快速判断类别轴在哪。希望这篇能帮你形成自己的判断流程。
阅读完成 · 觉得有帮助？
