JiahuiYu / slimmable_networks

Slimmable Networks, AutoSlim, and Beyond, ICLR 2019, and ICCV 2019
Other
914 stars 131 forks source link

fix up the problem in model_profiling when fed 'Modulelist' #59

Open 13015517713 opened 1 year ago

13015517713 commented 1 year ago

Problem Introduction:   Module.register_forward_hook function adds hooks to each module (with forward function) without ModuleList. So ModuleList module does not record the profile including macs and params. When traversing deep-first module.children fcuntion does not calculate the profile when this module is a instance of ModuleList.

Problem Reproduction:   You can try to modify "self.blocks" in model.py to the instance of ModuleList.

def __init__(self):
  self.blocks = nn.ModuleList()
def forward(self, x):
  for block in self.blocks:
    x = block(x)

Solution:   I handle ModuleList specially, recording to each non-ModuleList submodule.