How do I recover the 3D structure of a layer after a Fully Connected layer (or a flatten layer)?

I want to apply 1 fully connected layer first and then start apply convolutions. But I noticed that the 3D structure might be lost. Is there a way to recover that 3D structure and apply 2d convolutions or am I doomed to only apply 1d convolutions thereafter?

Cross posted:


You can unflatten after the fully connected layer:

original_size = input.size()
flat_input = input.view(original_size[0], -1)
flat_ouput = linear(flat_input)
output = flat_output.view(original_size)

I want to build an unflatten layer. How would one do that?

e.g. of the flatten layer:

#mini class to add a flatten layer to the ordered dictionary
class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0),-1)

from your code I assume the only way is to do this:

#mini class to add an unflatten layer to the ordered dictionary
class UnFlatten(nn.Module):
    def forward(self, flat_ouput, original_size):
        return flat_ouput.view(original_size)