How to get layer index by name in `nn.Sequential(...)`


(leh) #1

Hi guys, I want to use a CNN as a feature extractor. When defining a neural net with nn.Sequential, for example

self.features = nn.Sequential(OrderedDict({
        'conv_1': nn.Conv2d(1, 10, kernel_size=5),
        'conv_2': nn.Conv2d(10, 20, kernel_size=5),
        'dropout': nn.Dropout2d(),
        'linear_1': nn.Linear(320, 50),
        'linear_2': nn.Linear(50, 10)
}))

I wonder if there is any way I can get a layer’s index by its name. This way, it would be easier for me to extract feature. Imagine what I want in the code below:

indx = original_model.features.get_index_by_name('conv_1')
feature = original_model.features[:indx](x)

A more general question would be “how to extract features at specific layers” in a pretrained model defined with nn.Sequential. I hope I make it clear. Hope you guys can help me, thank you!


#2

Your get_index_by_name method could implement something like this small hack:

list(dict(features.named_children()).keys()).index('conv_2')
> 1

(leh) #3

Thank you #ptrblck. It works perfectly!