How can we provide the output of a Linear layer to a Conv2D in PyTorch?

Question

I am building an Autoencoder where I need to encode an image into a latent representation of length 100. I am using the following architecture for my model.

        self.conv1 = nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size=3)
        self.conv2 = nn.Conv2d(in_channels=32,out_channels=64,kernel_size=3,stride=2)
        self.conv3 = nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,stride=2)

        self.linear = nn.Linear(in_features=128*30*30,out_features=100)

        self.conv1_transpose = nn.ConvTranspose2d(in_channels=128,out_channels=64,kernel_size=3,stride=2,output_padding=1)
        self.conv2_transpose = nn.ConvTranspose2d(in_channels=64,out_channels=32,kernel_size=3,stride=2,output_padding=1)
        self.conv3_transpose = nn.ConvTranspose2d(in_channels=32,out_channels=3,kernel_size=3,stride=1)  

Is there any way I could give my Linear layer's output to a Conv2D or a ConvTranspose2D layer so that I can reconstruct my image? The output is restored if I remove the Linear layer. I want to know how I can reconstruct my image keeping the Linear layer

Any help would be appreciated. Thanks!

Answer 1

You could use another linear layer:

self.linear2 = nn.Linear(in_features=100, out_features=128*30*30)

And then reshape the output into a 3D volume and pass it into your de-convolution layers.