I'm creating a CNN network without other frameworks such as PyTorch, Keras, Tensorflow, and so on.

During the forward pass, the Flatten layer reshapes the previous layer's activation. I know there are a lot of questions about it, but what should I do with the Flatten layer during back-propagation? Should I compute the derivative of $dA$ and reshape it for the next layer or just reshape $dA$ of the previous layer?


The Flatten layer has no learnable parameters in itself (the operation it performs is fully defined by construction); still, it has to propagate the gradient to the previous layers.

In general, the Flatten operation is well-posed, as whatever is the input shape you know what the output shape is.

When you backpropagate, you are supposed to do an "Unflatten", which maps a flattened tensor into a tensor of a given shape, and you know what that specific shape is from the forward pass, so it is also a well-posed operation.

More formally

Say you have Img1 in input of your Flatten layer

$$ \begin{pmatrix} f_{1,1}(x; w_{1,1}) & f_{1,2}(x; w_{1,2}) \\ f_{2,1}(x; w_{2,1}) & f_{2,2}(x; w_{2,2}) \end{pmatrix} $$

So, in the output you have

$$ \begin{pmatrix} f_{1,1}(x; w_{1,1}) & f_{1,2}(x; w_{1,2}) & f_{2,1}(x; w_{2,1}) & f_{2,2}(x; w_{2,2}) \end{pmatrix} $$

When you compute the gradient you have

$$ \frac{df_{i,j}(x; w_{i,j})}{dw_{i,j}} $$

and everything in the same position as in the forward pass, so the unflatten maps from the (1, 4) tensor to the (2, 2) tensor.

  • $\begingroup$ As I understood, I'll need to compute $dA$ in Flatten layer, too. Then reshape it to input shape of previous layer (either Pooling or Conv2D). $\endgroup$ – TheFnafException Sep 5 '20 at 20:26
  • $\begingroup$ The OP had provided an answer that seemed more a confirmation of what you had written, so I converted that "answer" to a comment, as you can see it above this one. If you have some time, please, address that comment by providing some feedback, which, if I understand correctly, should be something like "yes, exactly". Then I will delete this comment. $\endgroup$ – nbro Dec 16 '20 at 15:02

Not the answer you're looking for? Browse other questions tagged or ask your own question.