[人工智能]maxpooling的backward，和高维np.array花式索引（高级索引）的几种方法

近日写maxpooling的时候，发现一个看似很简单，实则很麻烦的问题，属实算是我这段时间碰到最麻烦的问题，断断续续做了几天实验，最后用一个勉强算优美的形式解决了它，记录一下。

先说一下maxpooling，它的forward过程很简单，例如kernel=(C,kernelsize=2,kernelsize=2)，那么将原图以stride=2，进行一个非线性运算max，得到的就是池化层输出，例如（N,C,28,28）将输出为（N,C,14,14），这很简单。

但是它的backward过程呢，遗憾的是，至少在csdn上，我没有找到正确的、高维度的maxpooling的backward实现。

他们互相继承的时候，基本上犯了这两个错误：

1. backward的梯度应当只计算max位置，其它位置为0

2. 只有一张图、一个通道的计算，而其实现方法无法扩展到多张图、多通道

仔细想想，backward在把(N,C,14,14)形状的梯度，映射回(N,C,28,28)的时候，是只有(N,C,1,1)的max值的位置上有梯度，还是整个(N,C,2,2)的块上都有梯度呢？

应当是只有最大值所在位置有梯度，也就是Input.grad(N,C,28,28)中，只有1/4的位置是有梯度的，其它地方的梯度应当是0。

MaxUnpool2d — PyTorch 1.9.1 documentation

What is the backward function for a max pooling layer? - PyTorch Forums

maxpooling类比于一个相同kernel、stride=kernelsize的卷积层，其计算方法从”卷积“替换为”求max+置零“这组非线性计算，那么很好理解，置零的部分不应当有梯度。实验也证明，全部backward的情形下，模型将会震荡。

这个问题看似很简单，那么怎么写呢？或者说，怎样在不使用for和if，额外增加时间复杂度的情况下，甚至是，不增加空间复杂度的情况下，实现这个backward呢？

也就是说，怎样在只用高级索引和切片的情形下，实现这个backward？

先说一点，np.max可以取多个维度(axis=tuple)，但是np.argmax只能取一个维度(axis=int)，这两个是不对称的，就算你用np.unrave_index，那么高维度的情形呢。

此外，还有一个更难受的问题，np的高级索引，当某个维度送进去一个一维列表时，它将取这个列表中的每个值作为索引，读array的这一维度的索引，广播到其它所有维度。

那么高级索引中送进去一个高维矩阵呢？高级索引中送进去多个矩阵呢？

也就是说这几种情形：

写法一：array[...,[0,1],...] == np.concatenate((array[...,0,...],array[...,1,...]), axis=cat_axis)

并引申到：

写法二：array[...,[[0],[1]],...] == np.concatenate((array[...,[0],...],array[...,[1],...]), axis=cat_axis)

这两种写法是不一样的，写法一相当于该维度分别取一个（因为是索引，所以此时维度降了一维），再在对应维度拼接（此时维度升了一维）。而写法二，从写法一出发，维度不变，再拼接一次，所以升了一次维。

再引申一次：

写法三：array[...,[0,1],[0,1],...] == np.concatenate((array[...,0,0,...],array(...,1,1,...)),axis=cat_axis)

这里拼接的维度没法写了是错的，但是，原理上是分别取对应位置的维度（[0,0]和[1,1]），再拼接，所以此时维度先降2，再升1，总体是降1维。

现在回到我一开始的问题，maxpooling的backward怎么实现？或者说，怎样在input.grad中，

1. 在kernel块(N,C,kernersize,kernersize)中，求出(N,C,1,1)的索引“柱子”？

2. 在原图梯度块(N,C,kernersize,kernersize)中，索引到(N,C,1,1)的值并赋值？

第一个问题很简单，对于每个梯度块，先reshape成（N,C,-1）(记为arg_kernel)，就可以用np.argmax(axis=-1)，求出一个(N,C,1)的索引块了，那么再np.unravel_index(arg_kernel, shape=(kernelsize,kernelsize))（记为unrav），就可以得到一个(2,N,C)大小的索引块了，这里第一个维度值是[y,x]，代表了最值在kernel块中的相对位置，此时取unrav[0]就是y，unrav[1]就是x。

也就是说，此时在input.grad图中（N,C,28,28），对于每一个（N,C,kernelsize,kernelsize)块，都有一个索引y(N,C,1)和x（N,C,1），其中y在范围[0,kernelsize]之间，指最大值在这个块中的相对坐标。

那么此时，最难的问题来了，相对坐标有了，怎样取这个input.grad中的值？

我想到的第一种办法，也是作为反面教材的办法，既然高级索引如此复杂，那么不用就好了，于是我加了两层循环，对image_n和image_channel加了两层循环，结果很显然，速度直接掉下来十倍。

第二种方法，也是一个反面教材，是我踩的深坑，我的模型指标始终上不去，捋了几遍我的代码，还排查出了python传引用的几个错误（后来在父类中做了隔断）！最后才发现，我仍然忽略了max pooling的特性，当backward时，只回传一个，其余的置零。我的方法是，既然高级索引很复杂，那我用布尔索引可不可以。于是我求每个kernel块上的最值，构成一个和input.grad相同大小的max图，比较input.val==max_img，获得一个布尔索引，再对input.grad取布尔索引即可。但是这种方法，每个块都有可能会取到重复的max值，并不能反映每个块取一个值回传的数学属性。

我原本以为，每个块有多个最大值的情况很少，最终将收敛到一个比较均值的水平，虽然会引入噪声但是不影响；但现实狠狠的惊到我了，最值的重复比例达到一半以上，反而是退化成了原始未作处理的情形。

这时，引入第三种方法，也是我认为对numpy高级索引认识比较深的一种方法：

高级索引，本质上还是对矩阵的运算，那么自然的，遵守numpy的广播原则，我总结为：

从右到左，从低到高，有1扩1，无1对齐。

np.array的矩阵，写入时，从低维开始填充，也就是从shape右侧开始，向左（高维）填充；

用户侧读取时，刚好相反，从高向低，从左到右，开始读取。

对于每一个维度，从右开始，若为1，那么将这一维度广播到n和另一个矩阵对应维度(长度为n)对其；当不为1时，需要对应维度的长度相等，否则无法广播，将报错。

那么，高级索引的填充方法应当是（默认grad=0）：

self.input.grad[...,y:y+self.kernelSize,x:x+self.kernelSize]\
    [np.arange(iN)[...,None,None,None], np.arange(ic)[None,...,None,None],
     unrav[0][...,None,None],unrav[1][...,None,None]] = 1

这里解释一下，先取一个(N,C,2,2)的梯度块，记为target：

self.input.grad[...,y:y+self.kernelSize,x:x+self.kernelSize]

再对target，每个维度，填入一个（应当是）(N,C,1,1)大小的高级索引矩阵，常见的高级索引是一维数组的情形，是我这种情形在广播机制下的特殊情况，因此，取到(N,C,1,1)的正确写法应当是：

target[np.arange(iN)[..., None, None, None], np.arange(ic)[None, ..., None, None],
 unrav[0][..., None, None], unrav[1][..., None, None]]

其中unrav是np.unravel_index取到的kernel块内求最值时获得的相对索引，unrav[0]的shape是(iN,ic)，也就是所有图像、所有通道，在这一kernel块的最值的相对索引，再在低维度扩2维，成为(N,C,1,1)，target的剩下三个高级索引矩阵原理相同，最终每个高级索引矩阵都成为一个满足numpy广播机制的四维(N,C,H,W)的高维矩阵。

综上，maxpooling的backward解决了，外加（应该是）搞透了高级索引的使用方法。本文的backward的核心实现就是一个argmax,一个unravel_index，再加一个高级索引，但是着实需要对高级索引有透彻的认识才可以在n**2时间复杂度下完成这个任务。如果有更低时间复杂度的解法，烦请不吝赐教。