I have a model mimicking the region proposal network(rpn) of faster r-cnn, and I wish to collect some of the intermediate feature maps before the head of the rpn. At present my code looks like:

 for index, child in enumerate(self.backbone.children()):
        images = child(images)
        if index in self.layer_outputs_to_keep:
            feature_maps.append(images)


Since the dataset I am training on is very large, I want to make sure that this is indeed the fastest possible way to do so. Any insights/tips appreciated.