Encoder generates NaN during VAE learning process

nmsanm: <p>I looked into it and this seems to be a common problem, but I could not solve the problem.</p>
<p>I am trying to take one of the GymEnv environments as a pixel and have VAE learn it. However, in the process, the VAE value is <code>NaN</code> and the learning does not seem to be working.</p>
<p>Here is minimum code that can confirm the occurrence of <code>NaN</code> problem(since it is too long code, I upload it to pastebin):</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://pastebin.com/RuswMTd0">
  <header class="source">
      <img src="https://pastebin.com/favicon.ico" class="site-icon" width="16" height="16">

      <a href="https://pastebin.com/RuswMTd0" target="_blank" rel="noopener nofollow ugc">Pastebin</a>
  </header>

  <article class="onebox-body">
    <img width="" height="" src="https://pastebin.com/i/facebook.png" class="thumbnail">

<h3><a href="https://pastebin.com/RuswMTd0" target="_blank" rel="noopener nofollow ugc">import sysimport numpy as npimport torchfrom torch import nnfrom torch.n -...</a></h3>

  <p>Pastebin.com is the number one paste tool since 2002. Pastebin is a website where you can store text online for a set period of time.</p>


  </article>

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

  <div style="clear: both"></div>
</aside>
<p>
I have examined the output results within VAE and found that <code>NaN</code> is occurring at Encoder. However, I am stumped as to how to resolve this issue.</p>
<p>This will result in a Tensor that has become <code>NaN</code> in the VAE encoder after several rounds of learning. Below is an example of the results when the problem was actually reproduced in my environment.</p>
<pre><code class="lang-auto">input isnan1:False
encoded_x isnan2:False
encoded_x isnan3:False
encoded_x isnan4:False
torch.Size([1000, 32])
input isnan1:False
encoded_x isnan2:False
encoded_x isnan3:False
encoded_x isnan4:False
torch.Size([1000, 32])
input isnan1:False
encoded_x isnan2:False
encoded_x isnan3:False
encoded_x isnan4:False
torch.Size([1000, 32])
input isnan1:False
encoded_x isnan2:False
encoded_x isnan3:False
encoded_x isnan4:False
torch.Size([1000, 32])
input isnan1:False
encoded_x isnan2:True
encoded_x isnan3:True
tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;CloneBackward0&gt;)
tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;ReluBackward0&gt;)
</code></pre>
<p><strong>Edit</strong></p>
<p>I forgot write about installed packages. Here is result of <code>pip freeze</code>:</p>
<pre><code class="lang-auto">absl-py==2.0.0
certifi==2023.7.22
cffi==1.16.0
charset-normalizer==3.3.0
cloudpickle==3.0.0
colorama==0.4.6
contourpy==1.1.1
cycler==0.12.1
filelock==3.12.4
fonttools==4.43.1
fsspec==2023.9.2
glcontext==2.5.0
glfw==1.12.0
gym==0.26.2
gym-notices==0.0.8
gym3==0.3.3
idna==3.4
imageio==2.31.5
imageio-ffmpeg==0.3.0
Jinja2==3.1.2
kiwisolver==1.4.5
MarkupSafe==2.1.3
matplotlib==3.8.0
moderngl==5.8.2
mpmath==1.3.0
mujoco==2.2.0
networkx==3.1
numpy==1.26.1
packaging==23.2
Pillow==10.1.0
procgen==0.10.7
pycparser==2.21
PyOpenGL==3.1.7
pyparsing==3.1.1
python-dateutil==2.8.2
requests==2.31.0
six==1.16.0
sympy==1.12
tensordict==0.2.0
torch==2.1.0
torchrl==0.2.0
torchvision==0.16.0
tqdm==4.66.1
typing_extensions==4.8.0
urllib3==2.0.6
</code></pre>

J_Johnson: <ol>
<li>Have you tried substituting in a more standard loss function to see if that’s the issue?</li>
<li>You might try clipping the objective function. See here: <a href="https://spinningup.openai.com/en/latest/algorithms/ppo.html" class="inline-onebox" rel="noopener nofollow ugc">Proximal Policy Optimization — Spinning Up documentation</a></li>
<li>You might also try gradient clipping. See here: <a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.utils.clip_grad_norm_ — PyTorch 2.1 documentation</a></li>
</ol>
<p>For example:</p>
<pre><code class="lang-auto">loss.backward()

torch.nn.utils.clip_grad_norm_(model.parameters(), 5.)

optimizer.step()

</code></pre>
<p>Lastly, you might try lowering the learning rate of your optimizer.</p>

nmsanm: <p>Thank you for reply.<br>
I let you know about the points that I have been able to confirm.</p>
<ol>
<li>
<p>I changed loss function to BCE version and Gaussian loss version, but VAE’s Encoder output <code>NaN</code> in training phase.</p>
</li>
<li>
<p>I tried gradient clipping but VAE output <code>NaN</code> same as before.</p>
</li>
<li>
<p>Try lower learning rate (10^-4 to 10^-6) though, the result does not change from <code>NaN</code>. It seems to take longer training steps to output <code>NaN</code> by lowering learning rate, but it eventually output <code>NaN</code>.</p>
</li>
</ol>

J_Johnson: <p>Thanks for trying. That eliminates a few possibilities.</p>
<p>Was looking over the code again, and where are you scaling the image values from 0-255 to 0-1? I see you imported <a href="https://pytorch.org/rl/reference/generated/torchrl.envs.transforms.ToTensorImage.html" rel="noopener nofollow ugc">ToTensorImage</a> but didn’t use it.</p>

nmsanm: <aside class="quote no-group" data-username="J_Johnson" data-post="4" data-topic="190838">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/j_johnson/48/55494_2.png" class="avatar"> J_Johnson:</div>
<blockquote>
<p>Was looking over the code again, and where are you scaling the image values from 0-255 to 0-1? I see you imported <a href="https://pytorch.org/rl/reference/generated/torchrl.envs.transforms.ToTensorImage.html" rel="noopener nofollow ugc">ToTensorImage</a> but didn’t use it.</p>
</blockquote>
</aside>
<p>Thanks for your help! My problem seems to be solved by adding <code>ToTensorImage</code>!</p>
<p>Just out of my curiosity, why <code>ToTensorImage</code> prevent to appear a <code>NaN</code>? I think it’s related to the gradient computation, but I’m not really sure what’s actually happening.</p>

J_Johnson: <p>Images by default use int8 and the values fall between 0 and 255.</p>
<p>But for machine learning, you want the values scaled from 0 to 1. That class will rescale the images, and also permute the channels dim.</p>