PyTorch维度操作实战：从argmax到cumsum的dim参数精解

1. 张量维度基础：从括号嵌套看本质

刚接触PyTorch时，最让我头疼的就是理解dim参数。记得第一次看到torch.sum(dim=1)这样的代码时，完全不明白这个dim到底在控制什么。后来发现，理解dim的关键在于真正看懂张量的维度结构。

我们先用一个简单的三维张量来感受下维度结构：

python复制import torch
z = torch.ones(2,3,4)
print(z)

输出看起来像俄罗斯套娃：

code复制tensor([[[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]],

        [[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]]])

这里有个实用技巧：数最外层的中括号对数。最外层有2个中括号（用逗号分隔），这就是dim=0的大小；往里一层有3个中括号，是dim=1的大小；最内层有4个数字，是dim=2的大小。这种括号计数法是我教新手同事时必用的方法，比死记硬背维度数字直观多了。

2. dim参数的本质：控制变量法实践

经过多个项目的实践，我发现理解dim最有效的方法是"控制变量法"。想象dim参数就是在说："我要在这个维度上做操作，其他维度保持不变"。这就像做实验时只改变一个变量，观察它对结果的影响。

举个例子，当我们在二维矩阵上做sum(dim=0)时：

python复制a = torch.arange(6).view(2,3)
print(a.sum(dim=0))

输出是：

code复制tensor([3., 5., 7.])

这里dim=0意味着："我要在dim=0方向（行方向）上求和，但保持dim=1（列方向）不变"。所以计算时是把每一列的数字相加：第一列0+3=3，第二列1+4=5，第三列2+5=7。

3. 常用函数的dim行为详解

3.1 torch.argmax的维度魔法

argmax是我在图像分类任务中最常用的函数之一。它的dim参数决定了在哪个维度上寻找最大值索引。来看个实际案例：

python复制scores = torch.tensor