Issue with saving pruned model

[psvishnu1]

Hello,

I’m currently evaluating your pruning method and encountered an issue while trying to save the complete pruned model (not just the state_dict). My goal is to extend the pruned model and fine-tune it for my application, but I’m running into difficulties saving the pruned version correctly.

I’ve attempted the following:

• Reusing the load_pruned_checkpoint function.
• Modifying torch.save in logging.py to save the complete model.

However, in both cases, the saved model appears to have the same number of parameters as the original model, which leads me to believe that both methods are still saving the original (unpruned) model, not the pruned version.

To clarify my current test setup: I modified dino.sh to target only the dino_vitb16 model, use the pets dataset, and focus on the highest pruning ratio (to save time). I’m running this setup on a Colab instance with an A100 GPU.

Could you please suggest the correct approach to save and work on the pruned model?

Thank you for your help!

Regards,
Vishnu

[Fly-Yun-Maker]

May I ask if you have successfully reproduced the results of the paper

[WalterSimoncini]

Hi @psvishnu1, sorry for the late reply.

Indeed, if you load the model with load_pruned_checkpoint and print out the model, you will see the original in_features and out_features values for the linear layers, as load_pruned_checkpoint simply sliced the weight matrices without updating those fields. But if you run

sum(p.numel() for p in model.parameters())

on the pruned model, you will see that only a fraction of the original parameters remain. Nonetheless, I've updated load_pruned_checkpoint to re-initialize the linear layers, so you should now see the pruned dimensionalities when you run print(model).