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