【Pytorch】张量的维度/轴/dim的理解

1 前言

  在深度学习中一旦涉及到三维或者更高维的数组（张量），Pytorch中各种涉及张量的函数操作都会通过参数axis/dim来设定操作的维度。很多张量的运算、神经网络的构建，都会涉及到轴，但到底取哪个轴，却不是那么容易把握。
  下面会针对轴/axis/dim，基于 Pytorch 的代码和实例，尝试去理清张量运算中轴/axis/dim的设定。
  以下内容为自学的记录，其中多有借鉴他人的博客，并在参考文献中一并给出引用。

2 轴的概念

  对于一个张量，它的shape有几维，就对应有几个轴，也就对应着张量的层级。

import torch
a = torch.Tensor([[1,2,3], [4,5,6]]) 
# a shape [2, 3]
b = torch.Tensor([[7,8,9], [10,11,12]])
# b shape [2, 3]
c = torch.Tensor([[[1,2,3], [4,5,6]], [[7,8,9], [10,11,12]]])
# b shape [2, 2, 3]

对于a对象来讲，axis/dim=0 对应 shape [2, 3] 中的2，axis/dim=1 对应 shape [2, 3] 中的3；对于b对象来讲 axis/dim=0 对应 shape [2，2, 3] 中的2，axis/dim=1 对应 shape [2, 2，3] 中的第二个2，axis/dim=3 对应shape [2, 2，3] 中的3，dim=-1就是最后一维。

3 轴的使用举例

3.1 torch.cat()

  在做张量的拼接操作时，axis/dim设定了哪个轴，那对应的轴在拼接之后张量数会发生变化。具体来讲，axis/dim设定了哪个轴，就会遍历这个轴去做运算，其他轴顺序不变。结合代码和下图解释，dim=0意为按照行对矩阵进行拼接，dim=1为按照列进行拼接。

>> torch.cat((a,b), dim=0)
tensor([[ 1.,  2.,  3.],
        [ 4.,  5.,  6.],
        [ 7.,  8.,  9.],
        [10., 11., 12.]])

>> torch.cat((a,b), dim=1)
tensor([[ 1.,  2.,  3.,  7.,  8.,  9.],
        [ 4.,  5.,  6., 10., 11., 12.]])

  注意：torch.cat会增加现有维度的值,可以理解为续接，但是不会增加新的维度。torch.stack会新加增加一个维度，可以理解为叠加。

  进一步的三维的矩阵拼接同二维一样，根据指定的轴进行拼接，不会改变维度。

import torch

tensor1 = torch.ones(2, 2, 2)
tensor2 = torch.zeros(2, 2, 2)
print(torch.cat([tensor2, tensor1], dim=0))
print(torch.cat([tensor2, tensor1], dim=0).shape)  # （4, 2, 2）

print(torch.cat([tensor2, tensor1], dim=1))
print(torch.cat([tensor2, tensor1], dim=1).shape)  # （2, 4, 2）

print(torch.cat([tensor2, tensor1], dim=2))
print(torch.cat([tensor2, tensor1], dim=2).shape)  # （2, 2, 4）

3.2 torch.stack()

  【Pytorch基础】torch.stack()函数解析

3.3 torch.sum()

  【Pytorch进阶】torch.sum()中dim的使用

3.4 torch.mean()

  【Pytorch基础】二维/三维情况下 torch.mean()函数使用

3.5 torch.softmax()

4 参考文献

[1]NumPy中axis的理解
[2]torch.cat 和 torch.stack的区别
[3]torch.stack() 与 torch.cat()