How to use a conv2d layer after a flatten?

Question

I am not familiar with Deep learning and Pytorch. And I want to know how to deal, in general with such a situation. So, I was wondering if I used a pretrained model (EfficientNet for example) if I want to change the _fc attribute and use conv2d in it, how can I recover a 2D structure? Because the pretrained model flattens it just before _fc.

for example, the pretrained model outputs a flattened feature vector of 1280 elements what I did is the following:

self.efficient_net._fc = nn.Sequential(
                nn.Linear(1280, 1225),
                nn.Unflatten(dim=1, unflattened_size=(1, 35, 35)),
                nn.Conv2d(1, 35, kernel_size=1),
                ...,
                )

I didn't have a specific height and width to recover in the 2D structure, so I assumed that h = w = some size and I use a linear layer whose output is equal to the square of the "some size". In the example above 35² = 1225. I am not sure if the unflatten is the correct way to do this. Then I added the conv2d. My code works but it doesn't give good results which probably means that the 2D structure I recovered does not capture any meaningful information. Can anyone enlighten me with general knowledge about how things are done in my situation, or give me some comments? Thank you!

Answer 1

To answer the question in the title, your enclosed method is a valid way to use 2d convs after a flattened feature vector. However, the bad results you experience could come from the structure of your model or from the way you train it. Regarding you last question, it is very hard to give you an advice without knowing your intentions in detail. Regardless, here are my two cents.

First of all, in the case you actually want to use this approach, to have a pretrained model and add your layers after one of its layers, you might want to keep the parameters of the original network intact at least until your newly initialized layers get trained properly. To achieve that, you need to use the gradients for updating only the parts you added as described in this comment. (You definitely should read the other comments there as well to get a better picture.)

Secondly, it might be worth reconsidering if you really want to add your layers after the very last layers of the pretrained network. Depending on your goals, using the output of some prior layers as the input to your layers might be more beneficial to you. (Just do not forget to keep the parameters of the pretrained model intact as I advised in the prior point.)

Lastly, the structure of your layers should also be reconsidered. 2d convs with a kernel size of 1x1 in this situation seems strange to my limited experience, but without knowing what you want to do, its hard to give any solid advice in this regard.

Therefore you might be better off splitting your question into smaller parts and work your way through them one by one.

(My reputation is not high enough otherwise I would have left a comment.)