How to save the pruned model?

[xavierdawn]

After calling model = make_prune_model(model), I saved the model using both torch.save(model.state_dict(), os.path.join(pruned_path, 'pytorch_model.bin')) and model.save_pretrained(pruned_path), but it seems like the saved model is still the original one. Is this my mistake?

[Qi-Le1]

Hi @xavierdawn. Thanks for the interest in our work!

In model = make_prune_model(model), what we actually do is replace the model’s original layers with customized layers. These customized layers are designed to apply specific pruning techniques during the forward pass. However, immediately after calling model = make_prune_model(model), the model should still retain its original weights.

Please let me know if you have any other questions!

[xavierdawn]

Hi @xavierdawn. Thanks for the interest in our work!

In model = make_prune_model(model), what we actually do is replace the model’s original layers with customized layers. These customized layers are designed to apply specific pruning techniques during the forward pass. However, immediately after calling model = make_prune_model(model), the model should still retain its original weights.

Please let me know if you have any other questions!

Thank you for your prompt reply. So, is there any way to preserve the pruned model? Because I would like to fine-tune the pruned model on a custom dataset. Thank you for your assistance.