Extract intermediate representation of MiDaS neural network in pytorch?

Viewed 110

Pytorch documentation provides a concise way to apply MiDaS monocular depth estimation network for depth extraction. But how should I modify their code to extract network representation at some intermediate layer? I know that I could download the model from github and modify forward function to return what I want, but I am interested in the simplest solution, leaving outer code as is.

I'm aware of subclassing the model class and writing my own forward function, like here, but I don't know how to access the class in the code. The model instance is created straight away with midas = torch.hub.load("intel-isl/MiDaS", model_type). Maybe an example of using a forward hook will be easier.

1 Answers

As you said, using a forward hook on a nn.Module is the easiest way to go about it. Consider the documentation: https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_forward_hook

Basically you just have to define a function that takes three inputs (module, input, output) and then does whatever you want with that data. To find at what Module you want to place that hook you obviously need to be familiar with the structure of the model. You can just print(midas) to get a pretty-printed representation of all the modules available. I just chose some random one, and used the print() function as a hook:

midas.pretrained.model.blocks[3].mlp.fc2.register_forward_hook(print)

This means whenever we call midas(some_input), the hook (print in this case) will be called with the corresponding arguments. Of course instead of print you can write a function that saves those files to e.g. a list that you can access from the outside, or write them to a file etc.

Related