Try to get the output of a layer by register_forward_pre_hook

Kinny_Fly: <p>I’m work on KeypointRCNN, and I try to get the outputs of the FPN(feature maps), so I use “register_forward_pre_hook” like this:</p>
<p><strong>for name, m in model.named_modules(): **<br>
**        if name ==‘rpn.head’:</strong><br>
**            m.register_forward_pre_hook(get_layer)**</p>
<p>Where “get_layer” is a function saving the rpn.head’s input,which is the output of FPN(I guess):</p>
<p><strong>def get_layer(module, x): **<br>
**        all_features_map.append(x)</strong></p>
<p>Thanks to the poor memory of my 1050Ti GPU, I can only run one image each time, so I try this:</p>
<p><strong>for run_images , target in dataset:</strong><br>
**      with torch.no_grad():**<br>
**            model.eval()**<br>
**            model.to(‘cuda’)**<br>
**            run_images = torch.unsqueeze(run_images, 0).to(torch.float32).to(device)**<br>
**            predictions = model(run_images) **<br>
**      …**</p>
<p>However! I find that the dim[0] of outputs I get from “get_layer”, are linearly increasing.<br>
The first image,I get 1280,which is 256*5, well, it’s nice, but the next time I get 2560, and 3840 and so forth.<br>
I guess it’s due to for some reason, when I use “for”, the forward() function will somehow invoke more times.<br>
So I try to wrap this whole operation mentioned before as a function “get_feat”, I find that by this way, the dim[0] of features are always 1280, which is nice I guess.</p>
<p><strong>for i in range(40):</strong><br>
**    f,p,g = get_feat(i)**<br>
**    feat_set_save.append(f)**<br>
**    pos_set_save.append(p)**<br>
**    gt_set_save.append(g)**<br>
**    torch.cuda.empty_cache()**</p>
<p>However! For some unknown reasons, this way may require more memory, I meet “CUDA out of memory” every time when I run about 25 images this function. Actually I think it’s not my GPU’s problem,for the reason that when I didn’t warp this operation as a function, I have already run 1000 images(Although the features are wrong).<br>
Can u give me some advice? Thank u so much!!!</p>

ptrblck: <p>I cannot reproduce the increasing feature shape using a simple code snippet:</p>
<pre><code class="lang-python">model = models.detection.keypointrcnn_resnet50_fpn()
model.eval()
x = [torch.randn(3, 224, 224), torch.randn(3, 256, 256)]

all_features = []
def get_layer(module, x):
    all_features.append(x)

for name, m in model.named_modules():
    if name == 'rpn.head':
        model.register_forward_pre_hook(get_layer)

out = model(x)
print([f[0][0].shape for f in all_features])
</code></pre>
<p>So unsure what’s causing the issue. Your current code is a bit hard to read and you can post code snippets by wrapping them into three backticks ```, which would make debugging easier.</p>
<p>Note that you are not detaching the activations before storing them in the list, which would thus store the entire computation graph and all intermediate activations with it unless you either <code>detach()</code> the tensor or wrap the forward pass into a <code>with torch.no_grad()</code> guard (which will disallow the gradient computation).</p>

Kinny_Fly: <p>Thank u for your help.<br>
The reason you didn’t meet the problem I mentioned may be that u didn’t run this model in more than one batchs.<br>
I test this morning, if I use “for image, target in dataset:” or "for i in range(100):"I will get exactly same features, like this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61e9a892761e350cc20b5777e9b7d801634688d9.png" data-download-href="https://discuss.pytorch.org/uploads/default/61e9a892761e350cc20b5777e9b7d801634688d9" title="Untitled"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61e9a892761e350cc20b5777e9b7d801634688d9_2_690x157.png" alt="Untitled" data-base62-sha1="dYaVo7bIEkd3G4qxmbHoykOrxt7" width="690" height="157" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61e9a892761e350cc20b5777e9b7d801634688d9_2_690x157.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61e9a892761e350cc20b5777e9b7d801634688d9_2_1035x235.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61e9a892761e350cc20b5777e9b7d801634688d9_2_1380x314.png 2x" data-dominant-color="29323B"><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">Untitled</span><span class="informations">1661×379 63.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><br>
So I just simply choose one of those same features.And I carefully check my code.<br>
To some extent, I have solved my problem , but I have no idea how this happens, can u explain to me if u are interested?Thanks!</p>

Kinny_Fly: <p>Thank u for your patience!<br>
If u are interested, here is my demonstration:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/0/f0324a0f1e453a792b9c731aea69f4f77bc073f8.png" data-download-href="https://discuss.pytorch.org/uploads/default/f0324a0f1e453a792b9c731aea69f4f77bc073f8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f0324a0f1e453a792b9c731aea69f4f77bc073f8_2_690x323.png" alt="image" data-base62-sha1="ygSkUhknd9QDGEbwMwfejPsu7Us" width="690" height="323" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f0324a0f1e453a792b9c731aea69f4f77bc073f8_2_690x323.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/0/f0324a0f1e453a792b9c731aea69f4f77bc073f8_2_1035x484.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/0/f0324a0f1e453a792b9c731aea69f4f77bc073f8.png 2x" data-dominant-color="222B35"><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">1260×590 98.7 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>

ptrblck: <p>It’s expected that <code>all_features_map</code> increases in its size, since you are appending to the <code>list</code>, isn’t it?<br>
I’m not sure I understand the issue correctly, so feel free to post an executable code snippet and explain a bit more what the issue is. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

Kinny_Fly: <p>OK! I made a stupid mistake.<br>
I just figured out what had happened.<br>
A simple test is below:</p>
<pre><code class="lang-auto">import torchvision
import torch

model = torchvision.models.detection.keypointrcnn_resnet50_fpn()
model.eval()

all_features = []

class test_dataset(object):
    def __init__(self):
        pass
    def __getitem__(self, idx):
        return torch.rand((3,224,224)), torch.rand((1))
    def __len__(self):
        return 5

dataset = test_dataset()

for image,target in dataset:
    image = torch.unsqueeze(image,0) # use one image due to the limit of memeory in pratical use
    all_features = []
    def get_layer(module, x):
        all_features.append(x)
    model.rpn.head.register_forward_pre_hook(get_layer) # repeatly register!
    with torch.no_grad():
        out = model(image)    
    print(len(all_features))
</code></pre>
<p>And it will get:<br>
1 2 3 4 5 6…<br>
It’s because:</p>
<pre><code class="lang-auto">model.rpn.head.register_forward_pre_hook(get_layer) # repeatly register!
</code></pre>
<p>This function is used repeatedly, so the “forward()” function will get another hook every time, resulting in the linearly increasing of the features.<br>
Thank u for your coding advice! <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=9" title=":grinning:" class="emoji" alt=":grinning:"></p>