Converted PyTorch model is significantly larger and slower than the original
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
I successfully converted a pre-trained PyTorch model to CoreML. The converted model produces the expected output, but the .mlpackage file is more than 45 times larger than the original PyTorch .pth.tar file, and prediction takes more than 7 times longer.
Here’s [my fork](https://github.com/iansampson/DBT-net/tree/coreml) of the model repo. My conversion script is trace.py, which also includes one composite operator (for atan2). The model code is found in aia_inter_new.py and aia_trans.py. The pre-trained model that I use for conversion is BEST_MODEL/aia_merge_dns300_conti.pth.tar. If it’s useful, I can also provide the converted .mlpackage and the traced .pt file.
I’ll be grateful for your help :). This *seems* to be a bug, and I’m not sure how to go about resolving it on my own.
## To Reproduce
Clone the repo and check out the coreml branch:
```
git clone https://github.com/iansampson/DBT-Net.git
cd DBT-Net
git checkout coreml
```
Assuming PyTorch and coremltools are already installed, install these additional dependencies:
```
pip install ptflops
pip install librosa
pip install pesq
```
Trace and convert the model:
```
python trace.py
```
The script traces the model with JIT (using a random input of a fixed size), converts the traced model to CoreML, and prints out prediction times for the original model, the traced model, and the CoreML model. Both the JIT trace and the CoreML model are saved in the coreml directory.
On my machine (a 2018 13" MacBook Pro), trace.py prints out the following values:
```
Prediction time for original model: 15.275246999999979
Prediction time for traced model: 15.296491000000003
Prediction time for CoreML model: 108.631306
```
And the file sizes of the respective models are:
```
aia_merge_dns300_conti.pth.tar: 11.8 MB
dbt-net_aia_merge_dns300.pt: 12.2 MB
dbt-net_aia_merge_dns300.mlpackage: 538.8 MB
```
## System environment:
- coremltools version: 5.2.0
- PyTorch version: 1.10.1
- macOS version: Monterey
Contributor guide
Research direction
Clone the linked DBT-Net repository, check out the coreml branch, install the listed dependencies, and run trace.py. Read trace.py together with aia_inter_new.py and aia_trans.py, then compare the reported prediction times and model sizes for the original, traced, and CoreML models. Done means explaining or isolating why conversion produces the 538.8 MB model and slower predictions.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- macos, python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100