Gradient Solving for Convolutional Neural Networks

In some studies on the interpretability of neural networks, the loss function is always used to calculate the gradient of the last layer of feature maps, such as the famous Grad CAM, so the understanding of convolutional neural networks cannot just stay in the package At the stage, we need to disassemble the black box that calculates the gradient.

As shown in the figure, suppose there is a feature map $A$ , after a$2\times$The convolution kernelKK$of\; 2$$After\; the\; K$ operation, a new feature map$O$ , and then flatten it through$MLP$ gets an output vector YYof length 2$Y.$ _
If you want to know the feature map$The\; contribution\; of\; each\; element\; of\; A$ to the final output needs to be calculated$Y$ vsAAThe partial derivative of each element in$A$$\frac{\partial Y}{\partial A}$.
Let's sort out the feature map $A$ gets outputYYThe process of$Y$
$$O=CONV(A)$$
$$Y=MLP\left(O\right)$$
Therefore, according to the chain derivation rule,$\frac{\partial Y}{\partial A}=\frac{\partial Y}{\partial O}\frac{\partial O}{\partial A}$.
To output $Y_1=68$ example,$Y_1=0\ast O_{11}+1\ast O_{12}+0\ast O_{21}+1\ast O_{22}$, so $\frac{\partial Y_1}{\partial O}=[0101]$
Let's calculate again$\frac{\partial O}{\partial A}=\left[\begin{array}{cccccc}\frac{\partial O_{11}}{\partial A_{11}}& \frac{\partial O_{11}}{\partial A_{12}}& \frac{\partial O_{11}}{\partial A_{13}}& \frac{\partial O_{11}}{\partial A_{21}}& \dots & \frac{\partial O_{11}}{\partial A_{33}}\\ \frac{\partial O_{12}}{\partial A_{11}}& \frac{\partial O_{12}}{\partial A_{12}}& \frac{\partial O_{12}}{\partial A_{13}}& \frac{\partial O_{12}}{\partial A_{21}}& \dots & \frac{\partial O_{12}}{\partial A_{33}}\\ \frac{\partial O_{21}}{\partial A_{11}}& \frac{\partial O_{21}}{\partial A_{12}}& \frac{\partial O_{21}}{\partial A_{13}}& \frac{\partial O_{21}}{\partial A_{21}}& \dots & \frac{\partial O_{21}}{\partial A_{33}}\\ \frac{\partial O_{22}}{\partial A_{11}}& \frac{\partial O_{22}}{\partial A_{12}}& \frac{\partial O_{22}}{\partial A_{13}}& \frac{\partial O_{22}}{\partial A_{21}}& \dots & \frac{\partial O_{22}}{\partial A_{33}}\end{array}\right]=C^{After\; T}$
finally integrates the results, the shape transformation and$A$ can be the same, that is,$\left[\begin{array}{ccc}0& 0& 1\\ 0& 2& 4\\ 0& 2& 3\end{array}\right]$.
The following is the code of the above calculation process, and it can be found that the calculation results are consistent with the derivation.

import torch
import torch.nn as nn

X = torch.tensor([[0, 1, 2],
                  [3, 4, 5],
                  [6, 7, 8]]).reshape(1, 1, 3, 3).float()
X.requires_grad = True
kernel = torch.tensor([[0, 1],
                       [2, 3]]).reshape(1, 1, 2, 2).float()
conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=2, bias=False)
conv.weight.data = kernel

fc = nn.Linear(in_features=4, out_features=2, bias=False)
fc.weight.data = torch.tensor([[0, 1, 0, 1],
                               [1, 0, 1, 1]]).float()
print(conv(X))
O = fc(torch.flatten(conv(X), start_dim=1))
print(O)

O[0][0].backward()

print(X.grad)