search

lucaslie / torchprune

A research library for pytorch-based neural network pruning, compression, and more.
https://people.csail.mit.edu/lucasl/
MIT License
161 stars 24 forks source link

export to onnx #3

Open jakc4103 opened 2 years ago

jakc4103 commented 2 years ago

Hi Thanks for sharing this awesome code!

I noticed some customize pytorch modules is used in the trained pytorch models. I'm wondering if one could export the fine-tuned model to onnx after, for example, applying ALDS methods? Or is there any plan to support this?

It would be great if the trained model can be exported to onnx for deployment.

lucaslie commented 2 years ago

Hi, Thanks for raising this issue. That's a great idea!

It should be really easy to implement since all the checkpoints are handled by the trainer here.

I will implement it once I have some time to take care of it - otherwise, I would be happy to take PRs as well.