identity_mask_propagation issue while trying to prune yolov5

[abuvaneswari]

Hello,
I am trying to prune yolov5 model and my env is as follows:
torch==1.8.1
torchvision==0.9.1
nncf==1.7.1
Python 3.6.8
No CUDA
While the sparsity compression works fine, filter pruning does not work. I get the following error message:

INFO:nncf:Computing filter importance scores and binary masks...
Traceback (most recent call last):
File "train.nncf.py", line 539, in
train(hyp, opt, device, tb_writer, wandb)
File "train.nncf.py", line 246, in train
compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/model_creation.py", line 124, in create_compressed_model
compression_ctrl = composite_builder.build_controller(compressed_model)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/composite_compression.py", line 94, in build_controller
return self._child_builders[0].build_controller(model)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/pruning/filter_pruning/algo.py", line 66, in build_controller
self.config)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/pruning/filter_pruning/algo.py", line 108, in init
self.set_pruning_rate(self.pruning_init)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/pruning/filter_pruning/algo.py", line 336, in set_pruning_rate
self._apply_masks()
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/pruning/filter_pruning/algo.py", line 537, in _apply_masks
MaskPropagationAlgorithm(graph, PT_PRUNING_OPERATOR_METATYPES).mask_propagation()
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/common/pruning/mask_propagation.py", line 62, in mask_propagation
cls.mask_propagation(node, self._graph)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/pruning/export_helpers.py", line 144, in mask_propagation
identity_mask_propagation(node, graph)
File "/home/tfs/venv_nncf/lib/python3.6/site-packages/nncf/common/pruning/mask_propagation.py", line 80, in identity_mask_propagation
assert len(input_masks) == 1
AssertionError

I will share a zip of the directory with files.

To reproduce this error:

1. please populate the coco directory with coco dataset (untarred). Alternatively, modify the data/coco.yaml file as per your training data location.
2. pip install -r requirements.txt
3. train.nncf.py --data data/coco.yaml --cfg yolov5s.yaml --weights yolov5s.pt --batch-size 64

[abuvaneswari]

yolov5_pruning_nncf_issue.tar.gz

[abuvaneswari]

Download the yolov5s.pt from here and place it in the main directory:
https://github.com/ultralytics/yolov5/releases/download/v3.0/yolov5s.pt

[vshampor]

@AlexanderDokuchaev @mkaglins

[AlexanderDokuchaev]

Hi @abuvaneswari
Thanks for raising this issue with detailed description and code to reproduce.
Yes, there is problem with working on yolov5. We will soon fix this one.

[KrishnaKishore123]

Hi @AlexanderDokuchaev
I have faced the same issue recently. Any update from you on this?

[MaximProshin]

@daniil-lyakhov FYI