How do I free system RAM when from_pixels=True in SyncDataCollector?

Mphald: <p>I wrote the following reinforcement learning code in PyTorch, TorchRL. This changes the <code>total_frames</code> and <code>frames_per_batch</code>, but the code is basically based on the PPO tutorial: <a href="https://colab.research.google.com/drive/14vqi2NmNXW9MHQwc-XDTz61_iuLUDWZ2?usp=sharing" rel="noopener nofollow ugc">Google colab</a></p>
<p>In addition, a 96x96x3 image was obtained from the InvertedDoublePendulum in order to obtain visual data of the environment in the training.</p>
<pre data-code-wrap="python"><code class="lang-python">base_env = GymEnv("InvertedDoublePendulum-v4", device=device, frame_skip=frame_skip, from_pixels=True, pixels_only=False)
env = TransformedEnv(
    base_env,
    Compose(
        ObservationNorm(in_keys=["observation"]),
        PermuteTransform((-1, -3, -2), in_keys=["pixels"]),
        Resize(96, 96),
        PermuteTransform((-2, -1, -3), in_keys=["pixels"]),
        ToTensorImage(),
        DoubleToFloat(),
        StepCounter(),
    ),
)
</code></pre>
<p>This programme eats up memory as soon as it runs and is forced to terminate. And this eats up even 50 GB of RAM. I suspect this is because <code>SyncDataCollector</code> is dealing with pixel data.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/2/c2baef583be1d89b9533d84f41c7a1b59557acbe.png" alt="image" data-base62-sha1="rMF8oYofOYCLHgh5vEiPhoxpxKe" width="595" height="335"></p>
<p>However, image data is essential for me, so removing the <code>from_pixels</code> option is unthinkable. The <code>total_frames</code> is also necessary because if the total number of frames is reduced, the learning process is not completely finished. Changing the image size will sooner or later eat up RAM.</p>
<p>I need to save or free the RAM somehow. Do you have any ideas?</p>

vmoens: <aside class="quote no-group" data-username="Mphald" data-post="1" data-topic="209290">
<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/mphald/48/72622_2.png" class="avatar"> Mphald:</div>
<blockquote>
<pre><code class="lang-auto">        ObservationNorm(in_keys=["observation"]),
        PermuteTransform((-1, -3, -2), in_keys=["pixels"]),
        Resize(96, 96),
        PermuteTransform((-2, -1, -3), in_keys=["pixels"]),
        ToTensorImage(),
        DoubleToFloat(),
        StepCounter(),
</code></pre>
</blockquote>
</aside>
<p>My fist answer would be that you can reshuffle this and have a better perf by immediately resizing</p>
<pre><code class="lang-auto">        ToTensorImage(),
        Resize(96, 96),
        ObservationNorm(in_keys=["observation"]),
        DoubleToFloat(),
        StepCounter(),
</code></pre>
<p>Otherwise you are normalizing big images (which takes RAM) when you could do that after resize. Does that help?</p>

Mphald: <p>Thanks for the suggestion. I have tried it with Google Colab, but the problem of memory build up remains.</p>
<p>This is a cut-out of a Google Colab screen. The green indicator just to the left of <code>for i , tensordict...</code> indicates that this part is running.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/c/bc3c105bd9063be5add3169b92c405f046423401.png" data-download-href="https://discuss.pytorch.org/uploads/default/bc3c105bd9063be5add3169b92c405f046423401" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bc3c105bd9063be5add3169b92c405f046423401_2_690x115.png" alt="image" data-base62-sha1="qRcrWIH8KCnnWcNom0E3ToBq1hv" width="690" height="115" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bc3c105bd9063be5add3169b92c405f046423401_2_690x115.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bc3c105bd9063be5add3169b92c405f046423401_2_1035x172.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/c/bc3c105bd9063be5add3169b92c405f046423401_2_1380x230.png 2x" data-dominant-color="292A29"><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">1588×265 21.8 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>At the time of collecting data with this SyncDataCollector (<code>for i, tensordict_data...</code>) and at the <code>env.transform.init_stats</code>, memory usage increases and is accumulated as it is. As a result, the system RAM is wasted, as shown on the right, and the programme eventually stops.</p>
<hr>
<p>Not sure if this is relevant, here is the version of each package, for example. It is simply a <code>pip install</code> run on the Google Colab:</p>
<pre><code class="lang-auto">import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)

0.5.0 1.26.4 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] linux
</code></pre>

vmoens: <p>Can I ask why are you importing split_trajectories within the loop?<br>
It seems that loads the memory at every iteration. I don’t see any increase in RAM if I take it away.<br>
LMK if that makes sense</p>

Mphald: <p>That is the module I was going to use in an interations. I tried to remove that, but the memory still keeps growing.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/7/c71f42a258c5f755d5b9ef621dad327192df9a9d.png" data-download-href="https://discuss.pytorch.org/uploads/default/c71f42a258c5f755d5b9ef621dad327192df9a9d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c71f42a258c5f755d5b9ef621dad327192df9a9d_2_690x128.png" alt="image" data-base62-sha1="spvZtfJdMFYxByglf7TIzdDznQN" width="690" height="128" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c71f42a258c5f755d5b9ef621dad327192df9a9d_2_690x128.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c71f42a258c5f755d5b9ef621dad327192df9a9d_2_1035x192.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/7/c71f42a258c5f755d5b9ef621dad327192df9a9d_2_1380x256.png 2x" data-dominant-color="292A28"><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">1587×295 20.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>