https://github.com/akira4o4/tiny-convert
PyTorch -> ONNX （TorchScript）
https://github.com/akira4o4/tiny-convert
onnx python pytorch torchscript
Last synced: 5 months ago
JSON representation
PyTorch -> ONNX （TorchScript）
- Host: GitHub
- URL: https://github.com/akira4o4/tiny-convert
- Owner: akira4O4
- License: gpl-3.0
- Created: 2024-07-17T10:29:36.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2024-09-06T01:49:17.000Z (over 1 year ago)
- Last Synced: 2025-03-27T12:04:16.961Z (10 months ago)
- Topics: onnx, python, pytorch, torchscript
- Language: Python
- Homepage:
- Size: 19.5 KB
- Stars: 3
- Watchers: 1
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Tiny Model Convert Tool
---
## Introduction
This Repo is a tiny code for convert pytorch mode to other types model.
If have any questions or suggestions,please email me.
---
## Feat
- Support Windows,Linux,MacOS
- Convert **PyTorch** model to **ONNX** model (**Support dynamic shape**)
- Convert **PyTorch** model to **TorchScript** model
---
## How to install
```bash
cd
git clone https://github.com/akira4O4/tiny-convert.git
cd tiny-convert
pip install -r requirements.txt
```
---
## Load weight
```python
from src.utils import load_weight
model = model(...)
model_path = r''
load_weight(model, model_path)
```
---
## PyTorch → ONNX
If you need `dynamic shape`
```python
dynamic_axes = {
'images': {0: 'batch'},
'output': {0: 'batch'},
...
}
```
If you need `multi` inputs and outputs
```python
input_names = ['images1', 'images2', ...]
output_names = ['output1', 'output2', ...]
```
Run
```python
import torchvision.models as models
from src import VERSION
from src.export import Export
from src.utils import load_weight
if __name__ == '__main__':
# Create your model
net = models.resnet18(pretrained=True)
args = {
'model': net,
'mode': 'onnx',
'shape': (1, 3, 224, 224), # NCHW
'opset_version': 13,
'output': './',
'input_names': ['images'],
'output_names': ['output'],
'dynamic_axes': None,
'is_simplify': True,
}
export = Export(**args)
export.run()
```
---
## PyTorch → TorchScript
Run
```python
import torchvision.models as models
from src import VERSION
from src.export import Export
from src.utils import load_weight
if __name__ == '__main__':
net = models.resnet18(pretrained=True)
args = {
'model': net,
'mode': 'torchscript',
'shape': (1, 3, 224, 224), # NCHW
'output': './',
}
export = Export(**args)
export.run()
```