Mapping layer class with its parameters

I have a trained model and I want to do some surgery on it like merging layers into a fewer ones. For that, I need the following information for every layer in the model:

  1. layer class name - like nn.Conv2d, nn.BatchNorm2d or my custom classes.
  2. layer definition information - like what stride, what padding, bias true or false etc.
  3. layer parameters - like weights, bias, running mean & variance etc.

I have not been able to find a mechanism to get this consolidated information.

When I parse through the state dictionary, I only get a key and associated parameters. There is no way to find layer class/definition information from the key name. And when I use named_modules(), I get the layer information, but don’t get the key information to retrieve appropriate parameters from the state dictionary.

Appreciate your inputs, thanks!


When you go through named_modules, you have access to the parameters of each module. Why do you need to map it back to the state_dict?

@albanD, that is a great suggestion. Let me try that.