Can I access the inner layer outputs of DeepLab in pytorch?

Question

Using Pytorch, I am trying to implement a network that is using the pre=trained DeepLab ResNet-101. I found two possible methods for using this network:

this one

or

 torchvision.models.segmentation.deeplabv3_resnet101(
     pretrained=False, progress=True, num_classes=21, aux_loss=None, **kwargs)

However, I might not only need this network's output, but also several inside layers' outputs. Is there a way to access the inner layer outputs using one of these methods?

If not - Is it possible to manually copy the trained resnet's parameters so I can manually recreate it and add those outputs myself? (Hopefully the first option is possible so I won't need to do this)

Thanks!

Answer 1

You can achieve this without too much trouble using forward hooks.

The idea is to loop over the modules of your model, find the layers you're interested in, hook a callback function onto them. When called, those layers will trigger the hook. We will take advantage of this to save the intermediate outputs.

For example, let's say you want to get the outputs of layer classifier.0.convs.3.1:

layers = ['classifier.0.convs.3.1']
activations = {}

def forward_hook(name):
    def hook(module, x, y):
        activations[name] = y
    return hook

for name, module in model.named_modules():
    if name in layers:
        module.register_forward_hook(forward_hook(name))

*The closure around hook() made by forward_hook's scope is used to enclose the module's name which you wouldn't otherwise have access to at this point.

Everything is ready, we can call the model

>>> model = torchvision.models.segmentation.deeplabv3_resnet101(
        pretrained=True, progress=True, num_classes=21, aux_loss=None)

>>> model(torch.rand(16, 3, 100, 100))

And as expected, after inference, activations will have a new entry 'classifier.0.convs.3.1' which - in this case - will contain a tensor of shape (16, 256, 13, 13).

---

Not so long ago, I wrote an answer about a similar question which goes a little bit more in detail on how hooks can be used to inspect the intermediate output shapes.