Search code examples
pythondeep-learningpytorchtorchyolo

Get all layers including operators from torch in python


I want to get all layers like convolution etc... and operator like "add" in python but I really dont know how.

This is my code

import torch

# An instance of your model.
model = torch.jit.load('best5.torchscript')

# Print all the layers
for name, param in model.named_parameters():
    print(name, param.shape)

It works fine but problem is that I am able to see only layers like convolutional not all with operators. How can i fix that?

For example in this image I can see everything when model is loaded to netron enter image description here


Solution

  • replace named_parameters to named_modules and print module (in your case param) instead of param.shape