How to get the size of the output after ConvTranspose2d? - math

Question

I am trying to build a convolutional VAE. The inputs are images (tensors) of dimensions 320x800 and I should get an output size equal to the input size.
The problem is that I can't understand how the output size after ConvTranspose2d is calculated and so the final output size I obtain is different from the input one.
This is my code: by giving as inputs random tensors of size [3, 320, 800] I obtain tensors of size [3, 3, 64, 160]. Note that there are 3 channels because I am using rgb images. I hope you can help me.
import torch
import torch.nn as nn
import torch.nn.functional as F
init_channels = 64 # initial number of filters
image_channels = 3 # color channels
latent_dim = 100 # number of features to consider
# Define the model
# Initial immage size = 320 x 800
class ConvVAE(nn.Module):
def __init__(self):
super(ConvVAE,self).__init__()
# Encoder
# output size after this layer: 160 x 400
self.enc1 = nn.Conv2d(
in_channels=image_channels, out_channels=init_channels,
kernel_size=4, stride=2, padding=1
)
# output size after this layer: 80 x 200
self.enc2 = nn.Conv2d(
in_channels=init_channels, out_channels=init_channels*2,
kernel_size=4, stride=2, padding=1
)
# output size after this layer: 40 x 100
self.enc3 = nn.Conv2d(
in_channels=init_channels*2, out_channels=init_channels*4,
kernel_size=4, stride=2, padding=1
)
# output size after this layer: 20 x 50
self.enc4 = nn.Conv2d(
in_channels=init_channels*4, out_channels=init_channels*8,
kernel_size=4, stride=2, padding=1
)
# output size after this layer: 10 x 25
self.enc5 = nn.Conv2d(
in_channels=init_channels*8, out_channels=init_channels*16,
kernel_size=4, stride=2, padding=1
)
# output size after this layer: 4 x 10
self.enc6 = nn.Conv2d(
in_channels=init_channels*16, out_channels=1024,
kernel_size=(2,5), stride=2, padding=0
)
self.fc1 = nn.Linear(1024, 2048)
self.fc_mu = nn.Linear(2048, latent_dim)
self.fc_log_var = nn.Linear(2048, latent_dim)
self.fc2 = nn.Linear(latent_dim, 1024)
# Decoder
self.dec1 = nn.ConvTranspose2d(
in_channels=1024, out_channels=init_channels*16,
kernel_size=(2,5), stride=2, padding = 0
)
self.dec2 = nn.ConvTranspose2d(
in_channels=init_channels*16, out_channels=init_channels*8,
kernel_size=4, stride=2, padding = 1
)
self.dec3 = nn.ConvTranspose2d(
in_channels=init_channels*8, out_channels=init_channels*4,
kernel_size=4, stride=2, padding = 1
)
self.dec4 = nn.ConvTranspose2d(
in_channels=init_channels*4, out_channels=init_channels*2,
kernel_size=4, stride=2, padding = 1
)
self.dec5 = nn.ConvTranspose2d(
in_channels=init_channels*2, out_channels=init_channels,
kernel_size=4, stride=2, padding = 1
)
self.dec6 = nn.ConvTranspose2d(
in_channels=init_channels, out_channels=image_channels,
kernel_size=4, stride=2, padding = 1
)
def reparameterize(self, mu, log_var):
"""
:param mu: mean from the encoder's latent space
:param log_var: log variance from the encoder's latent space
"""
std = torch.exp(0.5*log_var) # standard deviation
eps = torch.randn_like(std) # `randn_like` as we need the same size
sample = mu + (eps * std) # sampling
return sample
def forward(self, x):
# encoding
x = F.relu(self.enc1(x))
x = F.relu(self.enc2(x))
x = F.relu(self.enc3(x))
x = F.relu(self.enc4(x))
x = F.relu(self.enc5(x))
x = F.relu(self.enc6(x))
batch, _, _, _ = x.shape
x = F.adaptive_avg_pool2d(x, 1).reshape(batch, -1)
hidden = self.fc1(x)
# get `mu` and `log_var`
mu = self.fc_mu(hidden)
log_var = self.fc_log_var(hidden)
# get the latent vector through reparameterization
z = self.reparameterize(mu, log_var)
z = self.fc2(z)
z = z.view(-1, 1024, 1, 1)
# decoding
x = F.relu(self.dec1(z))
x = F.relu(self.dec2(x))
x = F.relu(self.dec3(x))
x = F.relu(self.dec4(x))
x = F.relu(self.dec5(x))
reconstruction = torch.sigmoid(self.dec6(x))
return reconstruction, mu, log_var