Model speedup fails due to Attribute Error

Describe the bug: I am attempting to run this pruned mirror detection model [https://github.com/memgonzales/mirror-segmentation], but when I attempt to perform speedup_model(), I get the following error:

  File "C:\Users\pillai.k\AppData\Local\Programs\Python\Python39\lib\runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "C:\Users\pillai.k\AppData\Local\Programs\Python\Python39\lib\runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "C:\Users\pillai.k\PMDLite\mirror-segmentation-main\prune.py", line 178, in <module>
    main()
  File "C:\Users\pillai.k\PMDLite\mirror-segmentation-main\prune.py", line 127, in main
    ModelSpeedup(net, dummy, masks).speedup_model()
  File "C:\Users\pillai.k\MirrorEnv\lib\site-packages\nni\compression\speedup\model_speedup.py", line 429, in speedup_model
    self.fix_mask_conflict()
  File "C:\Users\pillai.k\MirrorEnv\lib\site-packages\nni\compression\speedup\model_speedup.py", line 243, in fix_mask_conflict
    fix_channel_mask_conflict(self.graph_module, self.masks)
  File "C:\Users\pillai.k\MirrorEnv\lib\site-packages\nni\compression\speedup\mask_conflict.py", line 229, in fix_channel_mask_conflict
    prune_axis = detect_mask_prune_dim(graph_module, masks)
  File "C:\Users\pillai.k\MirrorEnv\lib\site-packages\nni\compression\speedup\mask_conflict.py", line 400, in detect_mask_prune_dim
    sub_module = graph_module.get_submodule(layer_name)
  File "C:\Users\pillai.k\MirrorEnv\lib\site-packages\torch\nn\modules\module.py", line 686, in get_submodule
    raise AttributeError(mod._get_name() + " has no "
AttributeError: BFE_Module has no attribute `cbam`

This error occurs for the mask element edge_extract.cbam.ChannelGate.mlp.1. Given below is the model description for this module:

  (edge_extract): BFE_Module(
    (cbam): CBAM(
      (ChannelGate): ChannelGate(
        (mlp): Sequential(
          (0): Flatten()
          (1): Linear(in_features=56, out_features=3, bias=True)
          (2): ReLU(inplace=True)
          (3): Linear(in_features=3, out_features=56, bias=True)
        )
      )
      (SpatialGate): SpatialGate(
        (compress): ChannelPool()
        (spatial): BasicConv(
          (conv): Conv2d(1, 1, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3), bias=False)
          (bn): BatchNorm2d(1, eps=1e-05, momentum=0.01, affine=True, track_running_stats=True)
        )
      )
    )
    (edge_layer1): Sequential(
      (0): Conv2d(56, 28, kernel_size=(1, 1), stride=(1, 1))
      (1): BatchNorm2d(28, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
    (edge_layer2): Sequential(
      (0): Conv2d(56, 28, kernel_size=(3, 3), stride=(1, 1))
      (1): BatchNorm2d(28, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
    (edge_layer3): Sequential(
      (0): Conv2d(56, 28, kernel_size=(3, 3), stride=(1, 1), dilation=(2, 2))
      (1): BatchNorm2d(28, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
    (edge_layer4): Sequential(
      (0): Conv2d(56, 28, kernel_size=(3, 3), stride=(1, 1), dilation=(4, 4))
      (1): BatchNorm2d(28, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
  )

Environment:

NNI version:3.0
Training service (local|remote|pai|aml|etc):local
Python version: Python 3.9.4
PyTorch version: 2.1.2
Cpu or cuda version: cpu

Reproduce the problem

Code|Example: Model weights can be downloaded at https://drive.google.com/file/d/18zsqjK1aHVC4D8Ky530C--fwxdQylQ37/view?usp=sharing clone the above mentioned github and add the following to the main function of prune.py :


def main():
net = model.to(device)

# Load model weights and biases. Change the device ordinal as needed.
old_state_dict=torch.load(pruned_weights_path, map_location=device)
new_state_dict={}
for key in old_state_dict.keys():
    new_key=key.replace('module.','_nni_wrapper.')
    if('_mask' in key):
        tmp=new_key.split('.')
        tmp.insert(-1,"_nni_wrapper")
        new_key=(".".join(tmp))
    new_state_dict[new_key]=old_state_dict[key]
net.load_state_dict(new_state_dict)

pruner.unwrap_model()
fp=open('model_desc.txt','w')
print(net,file=fp)
fp.close()
from nni.compression.speedup import ModelSpeedup
ModelSpeedup(net, dummy, masks).speedup_model()
...rest of main()



- How to reproduce:
- set the path to pruned weights file and run `prune.py`

microsoft / nni

Model speedup fails due to Attribute Error #5794