一、dim简介

pytorch中对tensor的很多操作都涉及到dim（维度的设置），但是，我们总是搞不清楚每个维度代表什么，到底设置几维，比如sum（求和）、softmax、max（最大值）。ptorch中的dim类似于numpy纵的axis，这一次我们完全弄懂dim的使用

二、dim与方括号的关系

在进行指定维度的操作之前，我们先要明白dim与括号之间的关系，我们需要区分一下两个东西

问题：a1和a2是一样的吗？回答这个问题之前，我们先看下面：我们先创建一个矩阵

a = torch.tensor([[1, 2], [3, 4]])
print(a)

tensor([[1, 2],
        [3, 4]])

因为a是一个矩阵，所以a的左边有 2 个括号

括号之间是嵌套关系，代表了不同的维度。从左往右数，两个括号代表的维度分别是 0 和 1 ，在第 0 维遍历得到向量，在第 1 维遍历得到标量

同样地，对于 3 维 tensor

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)

tensor([[[3, 2],
         [1, 4]],

        [[5, 6],
         [7, 8]]])

则 3 个括号代表的维度从左往右分别为 0, 1, 2，在第 0 维遍历得到矩阵，在第 1 维遍历得到向量，在第 2 维遍历得到标量

更详细一点

所以，我们明白了：标量，向量，矩阵，三维张量之间的关系，三维张量里面包含了二维的矩阵，二维矩阵里面包含了一维的向量，一维向量里面包含了零维的标量。

我们再回到刚刚的问题，a1和a2肯定是不一样的，a1根本就不是三维的张量，a2才是三维张量的形式！

在某一维度求和（或者进行其他操作）就是对该维度中的元素进行求和。
对于矩阵 a

a = torch.tensor([[1, 2], [3, 4]])
print(a)

tensor([[1, 2],
        [3, 4]])

求 a 在第 0 维的和，因为第 0 维代表最外边的括号，括号中的元素为向量[1, 2]，[3, 4]，第 0 维的和就是第 0 维中的元素相加，也就是两个向量[1, 2]，[3, 4]相加，所以结果为

[1,2]+[3,4]=[4,6]

s = torch.sum(a, dim=0)
print(s)

tensor([4, 6])

可以看到，a 是 2 维矩阵，而相加的结果为 1 维向量，可以使用参数keepdim=True来保证维度数目不变。

s = torch.sum(a, dim=0, keepdim=True)
print(s)

tensor([[4, 6]])

在 a 的第 0 维求和，就是对第 0 维中的元素（向量）进行相加。同样的，对 a 第 1 维求和，就是对 a 第 1 维中的元素（标量）进行相加，a 的第 1 维元素为标量 1，2 和 3，4，则结果为

[1+2]=[3], [3+4]=[7]

s = torch.sum(a, dim=1)
print(s)

tensor([3, 7])

保持维度不变

s = torch.sum(a, dim=1, keepdim=True)
print(s)
tensor([[3],
        [7]])

对 3 维 tensor 的操作也是这样

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)

tensor([[[3, 2],
         [1, 4]],

        [[5, 6],
         [7, 8]]])

将 b 在第 0 维相加，第 0 维为最外层括号，最外层括号中的元素为矩阵[[3, 2], [1, 4]]和[[5, 6], [7, 8]]。在第 0 维求和，就是将第 0 维中的元素（矩阵）相加

s = torch.sum(b, dim=0)
print(s)
tensor([[ 8,  8],
        [ 8, 12]])

求 b 在第 1 维的和，就是将 b 第 1 维中的元素[3, 2]和[1, 4], [5, 6]和 [7, 8]相加，所以

[3,2]+[1,4]=[4,6],[5,6]+[7,8]=[12,14]

s = torch.sum(b, dim=0)
print(s)
tensor([[ 4,  6],
        [12, 14]])

则在 b 的第 2 维求和，就是对标量 3 和 2, 1 和 4, 5 和 6 , 7 和 8 求和

s = torch.sum(b, dim=2)
print(s)
tensor([[ 5,  5],
        [11, 15]])

除了求和，其他操作也是类似的，如求 b 在指定维度上的最大值

m = torch.max(b, dim=0)
print(m)

b 在第 0 维的最大值是第 0 维中的元素（两个矩阵[[3, 2], [1, 4]]和[[5, 6], [7, 8]]）的最大值，取矩阵对应位置最大值即可

torch.return_types.max(
values=tensor([[5, 6],
        [7, 8]]),
indices=tensor([[1, 1],
        [1, 1]]))

这里其实很好理解：在零维的时候，求和是两个矩阵对应位置的相加，而max操作也是两个矩阵对应位置的比较。

b 在第 1 维的最大值就是第 1 维元素（4 个（2对）向量）的最大值

m = torch.max(b, dim=1)
print(m)
torch.return_types.max(
values=tensor([[3, 4],
        [7, 8]]),
indices=tensor([[0, 1],
        [1, 1]]))

b 在第 0 维的最大值就是第 0 为元素（8 个（4 对）标量）的最大值

m = torch.max(b, dim=2)
print(m)
torch.return_types.max(
values=tensor([[3, 4],
        [6, 8]]),
indices=tensor([[0, 1],
        [1, 1]]))

在 tensor 的指定维度操作就是对指定维度包含的元素进行操作，如果想要保持结果的维度不变，设置参数keepdim=True即可。