Pytorch Tensor 维度操作的形象理解 Tensor.unsqueeze() Tensor.squeeze()

Posted 2022-03-11 西土城山羊卷

我们认为数组、矩阵、张量都是有形状的，假如有一个形状是 (2×2×3) 的张量 a，从左到右称为第0维、第1维、第2维…我可以使用 a[1] 取出第0维视角下的第1组数据，可以使用 a[:,:,0] 看到第2维视角下的第0组数据。

并且我可以使用 .unsqueeze 方法在 | 0 | 1 | 2| 这三个维度给出的4个任意空隙插入新的维度，如 b = a.unsqueeze(2) 则 b.size() = torch.size([2,2,1,3]) 。可以比较清晰的理解再最高的维度阔维的操作，这就是一些语言中 (如：Python MATLAB) 的广播机制对应的原理。但是在其他部位扩张维度是一件不好理解的事：

In [2]: import torch
In [3]: a = torch.tensor([[1, 2, 3], [4, 5, 6]])
In [4]: a
Out[4]: 
tensor([[1, 2, 3],
        [4, 5, 6]])
In [5]: a.size()
Out[5]: torch.Size([2, 3])
In [6]: b = a.unsqueeze(1)
In [7]: b
Out[7]: 
tensor([[[1, 2, 3]],

        [[4, 5, 6]]])
In [8]: b.size()
Out[8]: torch.Size([2, 1, 3])

对于一个 (2×3) 的张量，从第0维的一维视角看有两组数据 [1, 2, 3] [4, 5, 6]，从二维视角看就是一组平面数据，这个平面内实际上是有两个箱子的，一个箱子里放了 [1, 2, 3]，另一个箱子里放了 [4, 5, 6]，现在我给他在第一维视角上的箱子外又套了一层箱子。这就是阔维的思想。为了形象的理解张量形状的问题，用这种套箱思想，比线、平面、空间思想更加合理。这也与 Python 输出中的一层层方括号更加的契合。

In [9]: c = a.unsqueeze(2)
In [10]: c
Out[10]: 
tensor([[[1],
         [2],
         [3]],
        [[4],
         [5],
         [6]]])
In [11]: c.size()
Out[11]: torch.Size([2, 3, 1])

根据这种思想可以很好的理解上面的这种阔维做了什么。