Create_feature_extractor() does not find specified layers

Arturas_Druteika: <p>Hello, I want to create a cnn explainability class using <code>create_feature_extractor()</code> functionality. The problem I am facing is that I cannot reach certain layers in the models, for example (<strong>features.7.0</strong> is the last convolutional layer before avg pool):</p>
<pre><code class="lang-auto">from torchvision import models
from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names


model = models.efficientnet_v2_s(weights=models.EfficientNet_V2_S_Weights.DEFAULT)
model.eval()
return_nodes = ['features.7.0', 'classifier.1']
feature_extractor = create_feature_extractor(model, return_nodes=return_nodes)

</code></pre>
<p>This piece of code produces an error:</p>
<pre><code class="lang-auto">ValueError: node: 'features.7.0' is not present in model. Hint: use `get_graph_node_names` to make sure the `return_nodes` you specified are present. It may even be that you need to specify `train_return_nodes` and `eval_return_nodes` separately.
</code></pre>
<p>As the error points out I ran the <code>get_graph_node_names(model)</code> and I could not find this layer, even though I can do this <code>model.features[7][0]</code>, which gives me the last convolution.</p>
<p>If I pass the the return nodes as <code>return_nodes = ['features.7', 'classifier.1']</code> (i.e. without the .0 part) the feature extractor works, but the problem is that the features.7 is a Conv2dNormActivation which consists of 1 convolution, 1 batch norm and 1 activation function (as seen in the picture)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/e/5e9e2a55323bc5bacba3bd1001235281292c292f.png" data-download-href="https://discuss.pytorch.org/uploads/default/5e9e2a55323bc5bacba3bd1001235281292c292f" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/5/e/5e9e2a55323bc5bacba3bd1001235281292c292f.png" alt="image" data-base62-sha1="dv1KuNJ7MeFcPbjbarhlybzcG4v" width="690" height="90" data-dominant-color="283045"><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">721×95 3.15 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>Is there a possible solution for this problem so that I could get the output of this convolution operation from the feature etractor functionality?</p>

Alban_Grastien: <p>Hi folks.</p>
<p>I have a similar issue with layers whose name ends with <code>_&lt;number&gt;</code>.  There is a procedure that automatically renames the layers.  I am guessing it’s assuming the prefix was introduced during the traversal of the NN when the same module was encountered multiple times.<br>
Below is a small example that illustrates this:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torchvision.models.feature_extraction import get_graph_node_names
from torchvision.models.feature_extraction import create_feature_extractor

class M(nn.Module):
    def __init__(self):
        super(M, self).__init__()
        self.Lin_1 = nn.Linear( in_features=64, out_features=32, bias=True, )
        self.ACT_1 = nn.ReLU() # out_features=32
        self.Lin_2 = nn.Linear( in_features=32, out_features=16, bias=True, )
        self.ACT_2 = nn.ReLU()
        self.Lin_3 = nn.Linear( in_features=16, out_features=8, bias=True, )

    def forward(self, X):
        X = self.Lin_1(X)
        X = self.ACT_1(X)
        X = self.Lin_2(X)
        X = self.ACT_2(X)
        X = self.Lin_3(X)
        return X
    
if __name__ == '__main__':
    m = M()
    print(f'output shape = {m(torch.rand((64,))).shape}') # checking that the model is correct
    print(f'get_graph_node_names(m) = {get_graph_node_names(m)}')
    print(f'should return false: {"Lin" in get_graph_node_names(m)[0]}')
    print(f'should return true: {"Lin_3" in get_graph_node_names(m)[0]}')
    ce = create_feature_extractor(m, ['ACT_1'])
    output = ce(torch.rand((64,)))['ACT_1']
    print(f'out features should be 32: {output.shape}')
</code></pre>
<p>The output is:</p>
<pre><code class="lang-auto">output shape = torch.Size([8])
get_graph_node_names(m) = (['x', 'Lin', 'ACT', 'Lin_1', 'ACT_1', 'Lin_2'], ['x', 'Lin', 'ACT', 'Lin_1', 'ACT_1', 'Lin_2'])
should return false: True
should return true: False
out features should be 32: torch.Size([16])
</code></pre>
<p>My current solution is to rename the layers so that they have a different prefix; for instance, I simply add <code>x</code> at the end of their name (it’s annoying because I did not build the original NN).  Would there be a way to tell <code>get_graph_node_names</code> not to rename these layers?</p>

rxd: <p>Hey there Alban,<br>
the “_{0-9}+$” pattern is indeed removed automatically <a href="https://github.com/pytorch/vision/blob/518ee93dbd1469524040e1607a345fff90fa7fcd/torchvision/models/feature_extraction.py#L112C1-L118C60" rel="noopener nofollow ugc">here</a> (to compensate a potential post-fix added by torch.fx) and I don’t see any existing way to prevent this removal, other than your simple trick of renaming the layers.</p>
<p>As for the OP question, the issue comes from the definition of the <code>ConvNormActivation</code> <a href="https://github.com/pytorch/vision/blob/518ee93dbd1469524040e1607a345fff90fa7fcd/torchvision/ops/misc.py#L68" rel="noopener nofollow ugc">class</a>, which is lacking an explicit <code>__call__</code> function and is therefore treated as a black box when tracing the graph.<br>
With this change:</p>
<pre data-code-wrap="python"><code class="lang-python">class ConvNormActivation(torch.nn.Sequential):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int, ...]] = 3,
        stride: Union[int, Tuple[int, ...]] = 1,
        padding: Optional[Union[int, Tuple[int, ...], str]] = None,
        groups: int = 1,
        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
        dilation: Union[int, Tuple[int, ...]] = 1,
        inplace: Optional[bool] = True,
        bias: Optional[bool] = None,
        conv_layer: Callable[..., torch.nn.Module] = torch.nn.Conv2d,
    ) -&gt; None:

        if padding is None:
            if isinstance(kernel_size, int) and isinstance(dilation, int):
                padding = (kernel_size - 1) // 2 * dilation
            else:
                _conv_dim = len(kernel_size) if isinstance(kernel_size, Sequence) else len(dilation)
                kernel_size = _make_ntuple(kernel_size, _conv_dim)
                dilation = _make_ntuple(dilation, _conv_dim)
                padding = tuple((kernel_size[i] - 1) // 2 * dilation[i] for i in range(_conv_dim))
        if bias is None:
            bias = norm_layer is None

        layers = [
            conv_layer(
                in_channels,
                out_channels,
                kernel_size,
                stride,
                padding,
                dilation=dilation,
                groups=groups,
                bias=bias,
            )
        ]

        if norm_layer is not None:
            layers.append(norm_layer(out_channels))

        if activation_layer is not None:
            params = {} if inplace is None else {"inplace": inplace}
            layers.append(activation_layer(**params))
        super().__init__(*layers)
        _log_api_usage_once(self)
        self.out_channels = out_channels

        if self.__class__ == ConvNormActivation:
            warnings.warn(
                "Don't use ConvNormActivation directly, please use Conv2dNormActivation and Conv3dNormActivation instead."
            )

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)
</code></pre>
<p>access to <code>"features.7.0"</code> is possible. Hope it helps!</p>

Arturas_Druteika: <p>Hello. The way I have solved it is just to use good old hooks <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:" loading="lazy" width="20" height="20">. I couldn’t find any other workaround about this issue. Anyways, big thanks for the clarification on where the inner issue with the <strong>call</strong> method was.</p>