feature_extractor.py

class FeatureExtractor(nn.Module):
    def __init__(self, model, layer_names):
        super().__init__()
        self.model = model
        self.layer_names = layer_names
        self._features = defaultdict(list)
        
        layer_dict = dict([*self.model.named_modules()])
        for layer_name in layer_names:
            layer = layer_dict[layer_name]
            layer.register_forward_hook(self.save_outputs_hook(layer_name))

    def save_outputs_hook(self, layer_name):
        def fn(_, __, output): 
            self._features[layer_name] = output
        return fn

    def forward(self, x):
        _ = self.model(x)
        return self._features