VAE architecture and binary_cross_entropy

Omroth: <p>I’m trying haphazardly to implement a VAE on a two dimensional image.  I’d like to have the latent space also be 2d (because I read somewhere it was a good idea.)  This is my model:</p>
<pre><code class="lang-auto">class VAE(nn.Module):
    def __init__(self):
        super().__init__()
    
        # input shape = [128,128,1]

        self.el1 = nn.Conv2d(in_channels=n_input_channels,out_channels=n_latent_channels*2,kernel_size=4, padding=2, stride=4)
            
        # [32, 32, 64]

        self.el2 = nn.ReLU(inplace=True)

        # [32, 32, 64]

        self.el3 = nn.MaxPool2d(kernel_size=2, stride=2)


        # [16, 16, 64]
            
        self.dl1 = nn.ConvTranspose2d(in_channels=n_latent_channels, out_channels=n_latent_channels, kernel_size=[2,2], stride=[2,2])

        # [32, 32, 64]

        self.dl2 = nn.ConvTranspose2d(in_channels=n_latent_channels, out_channels=n_input_channels, kernel_size=[4,4], stride=[4,4])

        # [128, 128, 1]

    def reparameterise(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = std.data.new(std.size()).normal_()
            return eps.mul(std).add_(mu)
        else:
            return mu
    

    def encode(self, x):
        x = self.el1(x)
        x = self.el2(x)
        x = self.el3(x)

        latent = x

        mu = latent[:, :n_latent_channels, :, :]
        logvar = latent[:, n_latent_channels:, :, :]

        return mu, logvar

    def decode(self, z):
        z = self.dl1(z)
        z = self.dl2(z)

        return z

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterise(mu, logvar)

        x_hat = self.decode(z)
        return x_hat, mu, logvar

    def sample(self, n_samples):
        z = torch.randn((n_samples, n_latent_channels, 16, 16)).to(device)

        return self.decode(z)
</code></pre>
<p>And this is my loss function:</p>
<pre><code class="lang-auto">def loss_function(x_hat, x, mu, logvar):
    
    # reconstruction loss (pushing the points apart)
    BCE = nn.functional.binary_cross_entropy(
        x_hat, x, reduction='sum'
    )

    # KL divergence loss (the relative entropy between two distributions a multivariate gaussian and a normal)
    # (enforce a radius of 1 in each direction + pushing the means towards zero)
    KLD = 0.5 * torch.sum(logvar.exp() - logvar - 1 + mu.pow(2))

    return BCE + KLD # we can use a beta parameter here (BCE + beta * KLD)
</code></pre>
<ol>
<li>
<p>Is this architecture sane?</p>
</li>
<li>
<p>When I run the loss function I get</p>
</li>
</ol>
<p>“/aten/src/ATen/native/cuda/Loss.cu:115: operator(): block: [31,0,0], thread: [104,0,0] Assertion <code>input_val &gt;= zero &amp;&amp; input_val &lt;= one</code> failed”</p>
<p>which I assume means I haven’t normalised something during the encoder?</p>
<p>Any help would be greatly appreciated.</p>
<p>Thanks,<br>
Ian</p>

Omroth: <p>Well I normalised the input images to be in [0,1] and chucked a Sigmoid at the end of the decoder.  Which has solved this issue - does it sound like the right approach?</p>

ptrblck: <p>Yes, <code>nn.BCELoss</code> or its functional equivalent expect probabilities as the model output so using a <code>sigmoid</code> should fix the issue. For a better numerical stability you might want to remove the <code>sigmoid</code> and use <code>nn.BCEWithLogitsLoss</code>.</p>