Way to visualize which activations are stored during the forward pass?

username: Vedant_Roy
cooked: <p>Is there a way to view which activations are stored by a model during the forward pass? I’d like to figure out how to optimize the memory consumption of my model.</p>

username: Soumya_Kundu
cooked: <aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/johnmarktaylor91/torchlens">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/johnmarktaylor91/torchlens" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:666/500;"><img src="https://repository-images.githubusercontent.com/551190643/a758b584-7a74-4bec-b83a-7bc2dab612e7" class="thumbnail" width="666" height="500"></div>

<h3><a href="https://github.com/johnmarktaylor91/torchlens" target="_blank" rel="noopener nofollow ugc">GitHub - johnmarktaylor91/torchlens: Package for extracting and mapping the...</a></h3>

  <p>Package for extracting and mapping the results of every single tensor operation in a PyTorch model in one line of code.  - GitHub - johnmarktaylor91/torchlens: Package for extracting and mapping th...</p>


  </article>

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

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

<p>A simple example from the repo.</p>
<pre><code class="lang-auto">import torchlens as tl
import torch.nn as nn
import torch

class SimpleRecurrent(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(in_features=5, out_features=5)

    def forward(self, x):
        for r in range(4):
            x = self.fc(x)
            x = x + 1
            x = x * 2
        return x

x = torch.rand(6, 5)
simple_recurrent = SimpleRecurrent()
model_history = tl.log_forward_pass(simple_recurrent, x,
                                    layers_to_save='all',
                                    vis_opt='rolled')
print(model_history['linear_1_1:2'].tensor_contents)  # second pass of first linear layer

'''
tensor([[-0.0690, -1.3957, -0.3231, -0.1980,  0.7197],
        [-0.1083, -1.5051, -0.2570, -0.2024,  0.8248],
        [ 0.1031, -1.4315, -0.5999, -0.4017,  0.7580],
        [-0.0396, -1.3813, -0.3523, -0.2008,  0.6654],
        [ 0.0980, -1.4073, -0.5934, -0.3866,  0.7371],
        [-0.1106, -1.2909, -0.3393, -0.2439,  0.7345]])
'''
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/a/9afe5f82feae4db2c9d3a4654ea48cf86a6a8dc8.png" data-download-href="https://discuss.pytorch.org/uploads/default/9afe5f82feae4db2c9d3a4654ea48cf86a6a8dc8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/a/9afe5f82feae4db2c9d3a4654ea48cf86a6a8dc8_2_189x500.png" alt="image" data-base62-sha1="m78wnFv2T6o8mt9vd1rifvw5NlS" width="189" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/a/9afe5f82feae4db2c9d3a4654ea48cf86a6a8dc8_2_189x500.png, https://discuss.pytorch.org/uploads/default/original/3X/9/a/9afe5f82feae4db2c9d3a4654ea48cf86a6a8dc8.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/a/9afe5f82feae4db2c9d3a4654ea48cf86a6a8dc8.png 2x" data-dominant-color="E7E7E1"><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">190×502 26.9 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>