Loss goes to nan after a couple of epochs for a custom made VAE model

aquorio15: <p>Hi everyone,</p>
<p>I am training a VAE model which will take a list of numpy arrays and train a VAE model based on those arrays. My VAE model is inspired by the Transformer model as the input arrays is coming from a Image transformer. But after defining the model during training the loss is going to nan after the first epoch.<br>
My encoder model looks something like this</p>
<blockquote>
<p>class Encoder(nn.Module):<br>
def <strong>init</strong>(self, embedding_dim:int=768,<br>
mlp_size:int=512,<br>
m_size:int=256,<br>
dropout:int=0.1,<br>
latent_dims:int=128):<br>
super(Encoder, self).<strong>init</strong>()<br>
self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim)<br>
self.fc1 = nn.Linear(embedding_dim, 512)<br>
self.fc2 = nn.Linear(512, 256)<br>
self.fc3 = nn.Linear(256, latent_dims)<br>
self.fc4 = nn.Linear(256, latent_dims)<br>
self.N = torch.distributions.Normal(0, 1)<br>
self.N.loc = self.N.loc<br>
self.N.scale = self.N.scale<br>
self.kl = 0</p>
<pre><code>def forward(self, x):
    x = self.layer_norm(x)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    mu =  self.fc3(x)
    sigma = torch.exp(self.fc4(x))
    z = mu + sigma*self.N.sample(mu.shape)
    self.kl = (sigma**2 + mu**2 - torch.log(sigma) - 1/2).sum()
    return z
</code></pre>
</blockquote>
<p>My decoder model is</p>
<blockquote>
<p>class Decoder(nn.Module):<br>
def <strong>init</strong>(self, embedding_dim:int=768,<br>
mlp_size:int=512,<br>
m_size:int=256,<br>
dropout:int=0.1,<br>
latent_dims:int=128):<br>
super(Decoder, self).<strong>init</strong>()<br>
self.fc1 = nn.Linear(latent_dims, 256)<br>
self.fc2 = nn.Linear(256, 512)<br>
self.fc3 = nn.Linear(512, embedding_dim)</p>
<pre><code>def forward(self, z):
    x = F.leaky_relu(self.fc1(z))
    x = F.leaky_relu(self.fc2(x))
    x = self.fc3(x)
    return x
</code></pre>
</blockquote>
<p>My entire VAE model</p>
<blockquote>
<p>class VariationalAutoencoder(nn.Module):<br>
def <strong>init</strong>(self, embedding_dim:int=768,<br>
mlp_size:int=512,<br>
m_size:int=256,<br>
dropout:int=0.1,<br>
latent_dims:int=128):<br>
super(VariationalAutoencoder, self).<strong>init</strong>()<br>
self.encoder = Encoder(embedding_dim=768,<br>
mlp_size=512,<br>
m_size=256,<br>
dropout=0.1,<br>
latent_dims=128)<br>
self.decoder = Decoder(embedding_dim=768, latent_dims=128)</p>
<pre><code>def forward(self, x):
    z = self.encoder(x)
    y = self.decoder(z)
    return z, y
</code></pre>
</blockquote>
<p>I tried initializing the model with random weights and i found its working fine. But when i am training the loss goes to nan</p>
<blockquote>
<p>vae = VariationalAutoencoder(latent_dims=128)<br>
optimizer = torch.optim.Adam(vae.parameters(),<br>
lr=10-6,<br>
weight_decay=0.1)<br>
num_epochs = 10<br>
for epoch in range(num_epochs):<br>
for i, x in enumerate(train_dataloader_custom):<br>
z, y = vae(x)<br>
loss = torch.nn.MSELoss()<br>
reconst_loss = loss(y, x)<br>
optimizer.zero_grad()<br>
reconst_loss.backward()<br>
optimizer.step()<br>
print(reconst_loss)</p>
</blockquote>
<p>Thanks for any kind of help</p>

ptrblck: <p>I would recommend checking some intermediate tensors since e.g. <code>torch.log(sigma)</code> could create an invalid output easily depending on the output of <code>self.fc4</code>:</p>
<pre><code class="lang-python">sigma = torch.exp(torch.tensor(-500.))
torch.log(sigma)
# tensor(-inf)
</code></pre>

aquorio15: <p>Hi, thank you for the reply i am getting a nan loss even if i am not using the reconstruction loss. If i use simple mse loss i am getting nan values after the first epoch. This is the output from the mse loss function</p>
<blockquote>
<p>tensor(0.1416, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)<br>
tensor(nan, grad_fn=)</p>
</blockquote>

ptrblck: <p>Try to narrow down which operation or layer creates the NaN or Inf values, which could eventually create the NaNs.<br>
Using <code>torch.autograd.set_detect_anomaly(True)</code> might help.</p>

aquorio15: <p>where shall i put torch.autograd.set_detect_anomaly(True) before model creation or before defining loss function?</p>

ptrblck: <p>It doesn’t matter as long as it’s set before the model execution (forward/backward passes).</p>