Pytoch modifies the hook source code to obtain per layer output parameters (with layer name)

I want to use hook When obtaining the output parameters of each layer .

call pytorch Medium API, It is found that there is no good method to print the layer name and parameters of each layer one-to-one

For this reason, it is modified hook Source code , Just call again .

---

import torch
import torch.utils.hooks as hooks

def get_output_param(module, datasets):
	output_param = {
    }
	def hook(module, input, output):
		name = list(module._forward_hooks.keys())[0]
		output_param[name] = output.data
	for name, layer in module.named_modules():
		if name != "":
			handle = hooks.RemovableHandle(layer._forward_hooks)
			layer._forward_hooks[f"{
      name}"] = hook
	module(datasets)
	return output_param

end

---