Training with custom layer using unfold + mm + fold - Weights optimization

1

I am running into an issue where I have a custom convolutional layer using nn.Unfold, which I get the same values comparing with a regular nn.Conv2d with the same weights initialized.
However, if I replace the layer when training a network, the weights seem not to be optimized as the regular nn.Conv2d do, causing the network to be unable to learn.

Is there any additional step to implement in order to backpropagate the error the same way as the nn.Conv2d does?

Thank you for your attention!

2

torch.unfold or nn.Unfold would propagate the gradients.
Could you post your custom implementation of the conv module, so that we could have a look?

1 Like
4

Thank you very much for your answer, you are correct!
It was a small mistake of mine, concerning the use of the same nn.BatchNorm2D, before and after nn.Conv2d in the model’s forward pass… Apparently, there was a conflict I didn’t know about, my bad!
Specifying each nn.BatchNorm2D solved the problem. Thanks!