How can I extract features from all the layer from loaded GNN model?

Vishu_Gupta: <p>Recently, I git cloned a GNN network and am trying to get the output from all the layers (as I am not sure what each layer are named).</p>
<p>There is a file in the GitHub repo that does the prediction by loading a pre-trained model with an example input i.e. <code>pretrained.py</code> and I am trying to add some additional code in it to extract outputs.</p>
<p>I have tried adding the following code just after the model.eval() in the above-mentioned file that was mentioned in other threads (which was used to extract intermediate layer output from loaded CNN models):</p>
<pre><code class="lang-auto">activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

for name, layer in model.named_modules():
    layer.register_forward_hook(get_activation(name))
</code></pre>
<p>But I am getting the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "alignn/pretrained.py", line 248, in &lt;module&gt;
    out_data = get_prediction(
  File "alignn/pretrained.py", line 221, in get_prediction
    model([g.to(device), lg.to(device)])
  File "/home/vgf3011/.virtualenvs/alignn/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1120, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/vgf3011/vishu/alignn/alignn/models/alignn.py", line 290, in forward
    x, y, z = alignn_layer(g, lg, x, y, z)
  File "/home/vgf3011/.virtualenvs/alignn/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1120, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/home/vgf3011/vishu/alignn/alignn/models/alignn.py", line 163, in forward
    x, m = self.node_update(g, x, y)
  File "/home/vgf3011/.virtualenvs/alignn/lib64/python3.8/site-packages/torch/nn/modules/module.py", line 1123, in _call_impl
    hook_result = hook(self, input, result)
  File "alignn/pretrained.py", line 168, in hook
    activation[name] = output.detach()
AttributeError: 'tuple' object has no attribute 'detach'
</code></pre>
<p>Is there anything that I am doing wrong or do I have to be more specific on things to extract output from each of the layers?</p>

ptrblck: <p>As given in the error message <code>output</code> is a <code>tuple</code>, so you can’t call <code>detach()</code> on it (it’s a tensor operator) and would need to unwrap the <code>tuple</code> first.</p>

Vishu_Gupta: <p>Can you please let me know how to unwrap the <code>tuple</code>?</p>

ptrblck: <p>You could either iterate it or assign the content to separate variables:</p>
<pre><code class="lang-python">a = ((torch.tensor(1), torch.tensor(2)))
print(type(a))
# &gt; &lt;class 'tuple'&gt;

# iterate it
for x in a:
    print(x)
    
# unwrap
x, y = a
</code></pre>

Vishu_Gupta: <p>Is this something that I have to do at the place where I have tuple as the variable type, like in-between the layers of the GNN architecture?</p>

ptrblck: <p>No, you would have to unwarp the <code>tuple</code> in the <code>get_activation</code> method, as it raises the error here:</p>
<pre><code class="lang-python">activation[name] = output.detach()
</code></pre>

Vishu_Gupta: <p>When I tried to print the output type by <code>type(output)</code> inside the <code>get_activation</code> function, there seems to be output with type <code>tuple</code> at variable intervals. I am not sure if it is ok to ignore them for now but does it make sense to unwrap wherever you have output as type <code>tuple</code> and then detach it like:</p>
<pre><code class="lang-auto">out1, out2 = output
activation[name] = out1.detach()
activation[name] = out2.detach()
</code></pre>

ptrblck: <aside class="quote no-group" data-username="Vishu_Gupta" data-post="7" data-topic="140480">
<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/vishu_gupta/48/13778_2.png" class="avatar"> Vishu_Gupta:</div>
<blockquote>
<p>there seems to be output with type <code>tuple</code> at variable intervals</p>
</blockquote>
</aside>
<p>This would mean that different layers return a variable number of outputs (sometimes a single output would be returned otherwise multiple outputs).</p>
<aside class="quote no-group" data-username="Vishu_Gupta" data-post="7" data-topic="140480">
<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/vishu_gupta/48/13778_2.png" class="avatar"> Vishu_Gupta:</div>
<blockquote>
<p>I am not sure if it is ok to ignore them for now but does it make sense to unwrap wherever you have output as type <code>tuple</code> and then detach it like</p>
</blockquote>
</aside>
<p>You should check, if <code>output</code> is indeed a <code>tuple</code> before unwrapping it as otherwise you might slice a tensor or would run into a <code>ValueError</code>. Also, in your current code you are overriding the <code>activation[name]</code> value, so you should consider using different names to store the tuple tensors in the <code>activation</code> dict.</p>

Vishu_Gupta: <p>Thank You. Also, is there any way to switch on/off the activations i.e. we get the output of the layer with/without passing through the activation function while using the <code>get_activation</code> function?</p>

ptrblck: <p><code>activations</code> refers to the layer outputs in this context, not a specific activation function such as <code>nn.ReLU</code> so it’s not possible to “switch it off” without manipulating the <code>return</code> statement in the layer’s <code>forward</code> method.</p>

Vishu_Gupta: <p>I am not sure how to frame this correctly, but is the current output that we are getting the output from the layers only or do they pass through functions like ReLU before getting printed out ?</p>

ptrblck: <p>It depends on the model definition and in particular how the <code>forward</code> method is implemented.<br>
In your code snippet you are using:</p>
<pre><code class="lang-python">for name, layer in model.named_modules():
    layer.register_forward_hook(get_activation(name))
</code></pre>
<p>to register the forward hook for each module.<br>
If the activation functions (e.g. <code>nn.ReLU()</code>) are defined as modules via <code>self.act1 = nn.ReLU()</code>, a separate forward hook will be used for it. However, if the activation functions are used via the functional API in the <code>forward</code> via <code>x = F.relu(x)</code>, then no forward hook will be registered for them. In both cases, if the activation function uses <code>inplace=True</code> it’ll be applied (and visible) in the previous layer’s output.</p>

Vishu_Gupta: <p>I would also like to know if it is possible to get the output of a certain variable inside the forward pass such as variable <code>z = x + self.layer1(x)</code>? Can it be incorporated into the above loop or do I have to get it via some other method?</p>

ptrblck: <p>You could return <code>z</code> additionally to the original output in the <code>forward</code> method or e.g. add it to a global <code>list</code> etc.</p>