Variational Autoencoder with custom latent vector dimensions - target and input size mismatch

username: verance
cooked: <p>Hello, I am trying to build a simple autoencoder for images like these</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/6/e627981dfbd1e6a97514f5c19c819de0555b278a.png" data-download-href="https://discuss.pytorch.org/uploads/default/e627981dfbd1e6a97514f5c19c819de0555b278a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/6/e627981dfbd1e6a97514f5c19c819de0555b278a_2_690x462.png" alt="image" data-base62-sha1="wQ2EcsgDhBLx4CfTUsaLfF2urfI" width="690" height="462" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/6/e627981dfbd1e6a97514f5c19c819de0555b278a_2_690x462.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/6/e627981dfbd1e6a97514f5c19c819de0555b278a_2_1035x693.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/6/e627981dfbd1e6a97514f5c19c819de0555b278a.png 2x" data-dominant-color="676767"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1182×793 92.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The image size is 128x128.</p>
<p>I want my latent vectors to have 256 data points and I am struggling to build my NN for this purpose.</p>
<pre><code class="lang-auto">class ConvVAE(nn.Module):
    def __init__(self):
        super(ConvVAE, self).__init__()
 
        # encoder
        self.enc1 = nn.Conv2d(
            in_channels=image_channels, out_channels=init_channels, kernel_size=kernel_size, 
            stride=2, padding=1
        )
        self.enc2 = nn.Conv2d(
            in_channels=init_channels, out_channels=init_channels*2, kernel_size=kernel_size, 
            stride=2, padding=1
        )
        self.enc3 = nn.Conv2d(
            in_channels=init_channels*2, out_channels=init_channels*4, kernel_size=kernel_size, 
            stride=2, padding=1
        )
        self.enc4 = nn.Conv2d(
            in_channels=init_channels*4, out_channels=64, kernel_size=kernel_size, 
            stride=2, padding=0
        )
        # fully connected layers for learning representations
        self.fc1 = nn.Linear(64, 128)
        self.fc_mu = nn.Linear(128, latent_dim)
        self.fc_log_var = nn.Linear(128, latent_dim)
        self.fc2 = nn.Linear(latent_dim, 64)
        # decoder 
        self.dec1 = nn.ConvTranspose2d(
            in_channels=64, out_channels=init_channels*8, kernel_size=kernel_size, 
            stride=1, padding=0
        )
        self.dec2 = nn.ConvTranspose2d(
            in_channels=init_channels*8, out_channels=init_channels*4, kernel_size=kernel_size, 
            stride=2, padding=1
        )
        self.dec3 = nn.ConvTranspose2d(
            in_channels=init_channels*4, out_channels=init_channels*2, kernel_size=kernel_size, 
            stride=2, padding=1
        )
        self.dec4 = nn.ConvTranspose2d(
            in_channels=init_channels*2, out_channels=image_channels, kernel_size=kernel_size, 
            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))
        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, 64, 1, 1)
 
        # decoding
        x = F.relu(self.dec1(z))
        x = F.relu(self.dec2(x))
        x = F.relu(self.dec3(x))
        reconstruction = torch.sigmoid(self.dec4(x))
        return reconstruction, mu, log_var

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# initialize the model
model = ConvVAE().to(device)
# set the learning parameters
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.BCELoss(reduction='sum')

def final_loss(bce_loss, mu, logvar):
    """
    This function will add the reconstruction loss (BCELoss) and the 
    KL-Divergence.
    KL-Divergence = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    :param bce_loss: recontruction loss
    :param mu: the mean from the latent vector
    :param logvar: log variance from the latent vector
    """
    BCE = bce_loss 
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

def train(model, dataloader, dataset, device, optimizer, criterion):
    model.train()
    running_loss = 0.0
    counter = 0
    for i, data in tqdm(enumerate(dataloader), total=int(len(dataset)/dataloader.batch_size)):
        counter += 1
        data = data[0]
        data = data.to(device)
        optimizer.zero_grad()
        reconstruction, mu, logvar = model(data)
        bce_loss = criterion(reconstruction, data)
        loss = final_loss(bce_loss, mu, logvar)
        loss.backward()
        running_loss += loss.item()
        optimizer.step()
    train_loss = running_loss / counter 
    return train_loss

</code></pre>
<p>This is my code and I am trying do change it so my latent vector has 256 data points, but no success. Ive been trying to play with <code>latent_dim</code> parameter but I cannot make it work due to a dimensions problem.</p>
<blockquote>
<p>Using a target size (torch.Size([64, 1, 128, 128])) that is different to the input size (torch.Size([64, 1, 32, 32])) is deprecated. Please ensure they have the same size.</p>
</blockquote>
<p>Any idea what may be causing this? Can someone clear up how do I change my latent vector dimensions, which changes do I need to make to my NN architecture?</p>
<p>Kind regards</p>

username: SemiSimon
cooked: <p>I can’t tell at-a-glance what the problem would be, but i have a few suggestions.</p>
<p>First off, for tensor size mismatch problems, try to debug by peppering around your code<br>
<code>print("Tensor size at line x:", your_cool_tensor.size())</code> to keep track of tensor dimensions as they flow through the network. You’ll probably immediately notice where it goes wrong.</p>
<p>Second, i see you are using <code>ConvTranspose2d</code> alone for upsampling. I suggest against that, because you will face ugly checkerboard artifacts. Instead, change every <code>ConvTranspose2d</code> into an <code>Upsample</code> function (2x, bicubic) paired with a simple <code>Conv2D</code> applied on the upscaled tensor. kernel_size=3, stride=1 and padding=1, as to not change the size of the tensor, let the upscale take care of that. This would rid of the artifacts if you ever encounter them in the future.</p>
<p>Lastly, from the code it doesn’t appear that you are interested in using the latent representation other than for the loss nor to use the encoder/decoder separately (such as for using the decoder as a generator). If that’s the case, you might want to avoid the variational aspect entirely and end up with a much simpler network that would probably also help in narrowing down the cause of the problem.<br>
If you truly need a usable latent representation, you probably need to put the encoder and decoder in their own separate classes linked together by a master class, like in this tutorial:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/Jackson-Kang/Pytorch-VAE-tutorial/blob/master/01_Variational_AutoEncoder.ipynb">
  <header class="source">

      <a href="https://github.com/Jackson-Kang/Pytorch-VAE-tutorial/blob/master/01_Variational_AutoEncoder.ipynb" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/Jackson-Kang/Pytorch-VAE-tutorial/blob/master/01_Variational_AutoEncoder.ipynb" target="_blank" rel="noopener nofollow ugc">Jackson-Kang/Pytorch-VAE-tutorial/blob/master/01_Variational_AutoEncoder.ipynb</a></h4>


      <pre><code class="lang-ipynb">{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation of Variational AutoEncoder (VAE)\n",
    "\n",
    "     VAE from \"Auto-Encoding Variational Bayes\" (2014, D.P. Kingma et. al.)\n",
    "    \n",
    "     Kernel-author: Jackson Kang @ Deep-learning Lab. (Handong Global University, S.Korea)\n",
    "     \n",
    "     Author-email:  mskang@handong.edu\n",
    "\n",
    "     python and pytorch version: python=3.7, pytorch=1.3.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
</code></pre>



  This file has been truncated. <a href="https://github.com/Jackson-Kang/Pytorch-VAE-tutorial/blob/master/01_Variational_AutoEncoder.ipynb" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>