pytorch中维度dim的理解

PyTorch 中对 tensor 的很多操作如 sum、softmax 等都可以设置 dim 参数用来指定操作在哪一维进行。PyTorch 中的 dim 类似于 numpy 中的 axis，这篇文章来总结一下 PyTorch 中的 dim 操作。首先看一下这个图，图中给出了维度标号，注意区分正负，从左往右数，括号代表的维度分别是 0 和 1 和 2，从右往做为 -3 和 -2 和 -1。待会儿会用到。

图1

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

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

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

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

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

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

求 a 在第 0 维的和，因为第 0 维代表最外边的括号，括号中的元素为向量[1, 2]，[3, 4]，第 0 维的和就是第 0 维中的元素相加，也就是两个向量[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]])

同理的现在对dim=1进行操作

a = torch.tensor([[1,2],
                  [3,4]])
s = torch.sum(a,dim=1,keepdim=False)
print(s)
#输出 tensor([3, 7])

keepdim = True

a = torch.tensor([[1,2],
                  [3,4]])
s = torch.sum(a,dim=1,keepdim=True)
print(s)
# 输出 tensor([[3],
        [7]])

现在对三维张量进行操作

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 维中的元素（矩阵）相加

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
#print(b)
s = torch.sum(b,dim=0)
print(s)
# 输出
tensor([[ 8,  8],
        [ 8, 12]])
keepdim = True
#输出
 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]

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
#print(b)
s = torch.sum(b,dim=1)
print(s)
#输出 tensor([[ 4,  6],
              [12, 14]])
keepdim = True
#输出
tensor([[[ 4,  6]],
        [[12, 14]]])     

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

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
#print(b)
s = torch.sum(b,dim=2)
print(s)
#输出 tensor([[ 5,  5],
             [11, 15]])
keepdim = True
#输出 tensor([[[ 5],
         [ 5]],

        [[11],
         [15]]])    

现在再来看看其他dim有关的api

1.torch.max

在二维中

dim = 0

这个时候取的是矩阵的最大值以及下标即[[1,2],[3,4]]中的最大值和下标，那么应该是[3,4]

a = torch.tensor([[1,2],
                  [3,4]])
print(a)
print(torch.max(a,dim=0))
#输出
values=tensor([3, 4]),
indices=tensor([1, 1]))

dim = 1

这个时候取的是标量是在[1,2]和[3,4]中找到最大值

a = torch.tensor([[1,2],
                  [3,4]])
print(a)
print(torch.max(a,dim=1))
# 输出
values=tensor([2, 4]),
indices=tensor([1, 1]))

在三维中

dim = 0

则是对比这两个矩阵返回同位置最大的值

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)
print(torch.max(b,dim=0))
# 输出
tensor([[[3, 2],
         [1, 4]],

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

dim = 1

[3, 2], [1, 4]，[5, 6], [7, 8]这四个中返回最大值

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)
print(torch.max(b,dim=1))
# 输出
tensor([[[3, 2],
         [1, 4]],

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

dim = 2

对标量做比较

b = torch.tensor([[[3, 2], [1, 4]], [[5, 6], [7, 8]]])
print(b)
print(torch.max(b,dim=2))
# 输出
tensor([[[3, 2],
         [1, 4]],

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

那么同样的现在来思考下dim=-3，-2，-1的情况，图1中已经给出来了正负数维度对应的关系了，现在我们就取dim=-1看看

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

输出：

torch.return_types.max(
values=tensor([[3, 4],
        [6, 8]]),
indices=tensor([[0, 1],
        [1, 1]]))

发现了吗dim=-1和dim=2输出的是相同的，剩余的可自行验证。