Extract Features from models made with nn.ModuleList


I am using a 103 layer Tiramisu as mentioned here https://github.com/bfortuner/pytorch_tiramisu.
Now I would like to extract features from the bottleneck layer of the trained network, The authors have made use of nn.ModuleList for costructing the network. Any suggestions .

Rewriting forward function would be the easiest.
Hook would also works.

Okay thanks @checyuntc, I modified the forward function. Now, I would like to take the network use layers upto the bottleneck and add an additional class


I tried tweaking the forward pass and I think it was the easiest thing to do. However is there an nice way to truncate the model till the bottleneck layer, because I would now like to add an average pooling at the bottleneck layer to perform regression. Any ideas

here is an example of extract features from vgg with nn.Sequential. nn.ModuleList should be similar.

from torch import nn
from torchvision.models import vgg16

class Vgg16(nn.Module):
    def __init__(self):
        super(Vgg16, self).__init__()
        features = list(vgg16(pretrained = True).features)[:23]
        # the output of 3,8,15,22 layer is : relu1_2,relu2_2,relu3_3,relu4_3
        self.features = nn.ModuleList(features)
    def forward(self, x):
        results = []
        for ii,model in enumerate(self.features):
            x = model(x)
            if ii in {3,8,15,22}:
        return results

Maybe this recipe is useful:

Best regards