Search code examples
pythonpytorch

PyTorch get all layers of model


What's the easiest way to take a pytorch model and get a list of all the layers without any nn.Sequence groupings? For example, a better way to do this?

import pretrainedmodels

def unwrap_model(model):
    for i in children(model):
        if isinstance(i, nn.Sequential): unwrap_model(i)
        else: l.append(i)

model = pretrainedmodels.__dict__['xception'](num_classes=1000, pretrained='imagenet')
l = []
unwrap_model(model)            
            
print(l)
    

Solution

  • You can iterate over all modules of a model (including those inside each Sequential) with the modules() method. Here's a simple example:

    >>> model = nn.Sequential(nn.Linear(2, 2), 
                              nn.ReLU(),
                              nn.Sequential(nn.Linear(2, 1),
                              nn.Sigmoid()))
    
    >>> l = [module for module in model.modules() if not isinstance(module, nn.Sequential)]
    
    >>> l
    
    [Linear(in_features=2, out_features=2, bias=True),
     ReLU(),
     Linear(in_features=2, out_features=1, bias=True),
     Sigmoid()]