Fail to prune yolox

[Sugar929]

I try to prune yolox from https://github.com/Megvii-BaseDetection/YOLOX

pruner = tp.pruner.MagnitudePruner(
      model,
      example_inputs,
      importance=imp,
      iterative_steps=iterative_steps,
      ignored_layers=[],
      ch_sparsity_dict={model.backbone.backbone.dark2[0].conv:0.5},
      ch_sparsity= 0.0,
    )

Error:

  File "/home/miniconda3/lib/python3.8/site-packages/torch_pruning/importance.py", line 97, in __call__
    group_imp = torch.stack(group_imp, dim=0)
RuntimeError: stack expects each tensor to be equal size, but got [256] at entry 0 and [1024] at entry 2

I followed your advice in https://github.com/VainF/Torch-Pruning/issues/100 group_imp = torch.stack(group_imp, dim=0) ->group_imp = torch.cat(group_imp) But still error:

  File "/home/miniconda3/lib/python3.8/site-packages/torch_pruning/pruner/algorithms/metapruner.py", line 226, in prune_local
    pruning_idxs = imp_argsort[:(n_pruned//ch_groups)]
IndexError: dimension specified as 0 but tensor has no dimensions

[VainF]

Hi @Sugar929,

We are currently in the process of improving our detector pruning capabilities to include better support for popular networks like YOLO v1-v7. Unfortunately, due to the upcoming deadline for the ICCV conference, we may require a little more time to address this bug. We appreciate your patience and understanding as we work towards resolving this issue.

Besides, really thanks for your debug information. we will prioritize the support for yolox.

[Sugar929]

That's all right. Good luck with ICCV! And I have another question, why I moved the model to gpu, but after pruning(I used the torch_pruning==0.2.4) the model was on cpu. So when doing pruning, the model was moved back to cpu automatically?

[VainF]

Yes. in the old version, the model will be automatically moved to CPU. This behavior has been changed in the latest version.