can't speed up mobilenetv2 with slim

[twmht]

Environment:

• NNI version:2.0

Here are the weights

m0.25.zip

python amc_train --ckpt_path m0.25.pth --mask_path mask_m0.25.pth

would throw

[2021-01-19 14:33:39] INFO (nni.compression.pytorch.speedup.compressor/MainThread) start to speed up the model
[2021-01-19 14:33:39] INFO (nni.compression.pytorch.speedup.compressor/MainThread) fix the mask conflict of the interdependent layers
[2021-01-19 14:33:40] INFO (nni.compression.pytorch.utils.mask_conflict/MainThread) {'features.0.0': 8, 'features.1.conv.0': 8, 'features.1.conv.3': 1, 'features.2.conv.0': 24, 'features.2.conv.3': 24, 'features.2.conv.6': 1, 'features.3.conv.0': 36, 'features.3.conv.3': 36, 'features.3.conv.6': 1, 'features.4.conv.0': 36, 'features.4.conv.3': 36, 'features.4.conv.6': 1, 'features.5.conv.0': 48, 'features.5.conv.3': 48, 'features.5.conv.6': 1, 'features.6.conv.0': 48, 'features.6.conv.3': 48, 'features.6.conv.6': 1, 'features.7.conv.0': 48, 'features.7.conv.3': 48, 'features.7.conv.6': 1, 'features.8.conv.0': 96, 'features.8.conv.3': 96, 'features.8.conv.6': 1, 'features.9.conv.0': 96, 'features.9.conv.3': 96, 'features.9.conv.6': 1, 'features.10.conv.0': 96, 'features.10.conv.3': 96, 'features.10.conv.6': 1, 'features.11.conv.0': 96, 'features.11.conv.3': 96, 'features.11.conv.6': 1, 'features.12.conv.0': 144, 'features.12.conv.3': 144, 'features.12.conv.6': 1, 'features.13.conv.0': 144, 'features.13.conv.3': 144, 'features.13.conv.6': 1, 'features.14.conv.0': 144, 'features.14.conv.3': 144, 'features.14.conv.6': 1, 'features.15.conv.0': 240, 'features.15.conv.3': 240, 'features.15.conv.6': 1, 'features.16.conv.0': 240, 'features.16.conv.3': 240, 'features.16.conv.6': 1, 'features.17.conv.0': 240, 'features.17.conv.3': 240, 'features.17.conv.6': 1, 'features.18.0': 1}
[2021-01-19 14:33:40] WARNING (nni.compression.pytorch.utils.mask_conflict/MainThread) no multi-dimension masks found.
[2021-01-19 14:33:40] INFO (nni.compression.pytorch.utils.mask_conflict/MainThread) detected conv prune dim: 0
[2021-01-19 14:33:40] INFO (nni.compression.pytorch.speedup.compressor/MainThread) infer module masks...
Traceback (most recent call last):
  File "amc_train.py", line 215, in <module>
    net = get_model(args)  # for measure
  File "amc_train.py", line 84, in get_model
    ms.speedup_model()
  File "/home/shared/nfs/acer-share/nni/nni/compression/pytorch/speedup/compressor.py", line 183, in speedup_model
    self.infer_modules_masks()
  File "/home/shared/nfs/acer-share/nni/nni/compression/pytorch/speedup/compressor.py", line 140, in infer_modules_masks
    self.infer_module_mask(module_name, None, mask=mask)
  File "/home/shared/nfs/acer-share/nni/nni/compression/pytorch/speedup/compressor.py", line 120, in infer_module_mask
    self.infer_module_mask(_module_name, module_name, out_shape=input_cmask)
  File "/home/shared/nfs/acer-share/nni/nni/compression/pytorch/speedup/compressor.py", line 115, in infer_module_mask
    input_cmask = infer_from_outshape[m_type](module_masks, out_shape)
  File "/home/shared/nfs/acer-share/nni/nni/compression/pytorch/speedup/infer_shape.py", line 284, in <lambda>
    'Conv2d': lambda module_masks, mask: conv2d_outshape(module_masks, mask),
  File "/home/shared/nfs/acer-share/nni/nni/compression/pytorch/speedup/infer_shape.py", line 1001, in conv2d_outshape
    module_masks.output_mask.mask_index[1] == mask.mask_index[1])
AssertionError

any idea?

[mlbo]

I meet such issue too

[zheng-ningxin]

Thanks for the valuable feedback!
I encountered a little problem while reproducing your problem. I run the following cmd and get a size mismatch error from Pytorch.

python amc_train.py --mask_path mask_m0.25.pth --model_type mobilenetv2 --ckpt_path m0.25.pth
So I guess you have updated the amc_train.py or use your own mobilnetv2?
Could you please provide your amc_train.py so that I can reproduce this problem, thanks!

[mlbo]

This is my amc_train.py file, I do this successful when i use nni1.9
amc_train.zip

[scarlett2018]

@zheng-ningxin - would you mind share the latest status of this issue? thanks!

[zheng-ningxin]

Sincerely sorry for the late reply, we will fix this in v2.4. Please try again after this pr merged. Thanks!
#4008