pytorch中元素维度扩张和广播的个人理解

对于矩阵张量增加维度和广播的理解：

假设我们有一个2*4的矩阵

我们想要实现一个操作，将每个坐标分别计算它与其他坐标的相对位置差，即对于图上的坐标(0, 0),分别计算与其他4个坐标（包含本身）的差，于是我们得到(0, 0),(0, -1),(-1, 0),(-1, -1)

可以看到，计算一个坐标，我们得到的结果是坐标本身(0,0)信息量的四倍。如果我们对每个坐标都做这样的运算，代码设计将会是十分繁琐的，因此我们引入了pytorch中的维度扩张和广播机制如下：

首先，我们将矩阵通过python的列表表示：
[      [0, 0, 1, 1],

       [0, 1, 0, 1]      ]

记为fmp

接着，我们在最后一个维度后和最后一个维度之前插入一个维度

fmpa = fmp[:, ;, None] -> [2, 4, 1]

fmpb[:, None, :] -> [2, 1, 4]

接着我们进行相减的操作

relative = fmpa - fmpb

这时候，会自动引入pytorch的广播机制，即分别将fmpa和fmpb的形状扩充为[2, 4, 4]

即fmpa =

[      [[0, 0, 0, 0],

       [0, 0, 0, 0],

       [1, 1, 1, 1],

       [1, 1, 1, 1]],

    [[0, 0, 0, 0],

       [1, 1, 1, 1],

       [0, 0, 0, 0],

       [1, 1, 1, 1]      ]]

这里可能比较绕，我直接给出规律，即对于新的矩阵第二个维度的第一个矩阵的每一行相当于原来每个元素的x坐标重复了4次，第二行是原来的第二个元素的x坐标重复了4次，以此类推，也就是说第二个维度的第一个矩阵[:0:]保存了原本矩阵的4个元素的x坐标；于是，第二个维度的第二个矩阵[:1:]也就保存了原本矩阵的每个元素的y坐标

对于fmpb：

[      [[0, 0, 1, 1],

       [[0, 0, 1, 1],

       [[0, 0, 1, 1],

       [[0, 0, 1, 1]],

       [[0, 1, 0, 1],

       [[0, 1, 0, 1],

       [[0, 1, 0, 1],

       [[0, 1, 0, 1]]]

可以看到，从第二个维度开始，好像是fmpa转置了一样，也就是说，对于第二个维度的第一个矩阵的每一行，保存的是原来fmp矩阵的每一个元素的i坐标，第二个矩阵保存了fmp矩阵的每个元素的j坐标。

于是，[: 0:]保存的是i坐标，[:1 :]保存的是j坐标，只不过fmpb是在行方向递增原来矩阵的i或者j坐标，列方向进行重复，而fmpa是在行方向进行重复，列方向递增原来的i或者j坐标。因此，如果这个时候我们进行矩阵相减fmpa – fmpb，就实现了对于每个元素，计算其与其他元素（包含本身）的位置差。

如果还是感觉很抽象，可以结合b站up主@霹雳吧啦Wz的视频《使用Pytorch搭建Swin-Transformer网络》，跳转到60：00