Add rten-convert (#53)

Convert RTen model converter into an installable Python package.

Convert `tools/convert-onnx.py` into a Python package.

Fixes #45.

---------

Co-authored-by: Igor Yusupov <[email protected]>
Co-authored-by: Robert Knight <[email protected]>

.gitignore

@@ -2,6 +2,7 @@ __pycache__
 /dist
 /node_modules
 /target
+/rten-convert/rten_convert.egg-info
 
 # Optimized neural network models
 *.model

Makefile

@@ -1,5 +1,5 @@
 .PHONY: all
-all: src/schema_generated.rs tools/schema_generated.py
+all: src/schema_generated.rs rten-convert/rten_convert/schema_generated.py
 
 .PHONY: clean
 clean:
@@ -51,8 +51,8 @@ src/schema_generated.rs: src/schema.fbs
 	(echo "#![allow(clippy::all)]" && cat src/schema_generated.rs) > src/schema_generated.rs.tmp
 	mv src/schema_generated.rs.tmp src/schema_generated.rs
 
-tools/schema_generated.py: src/schema.fbs
-	flatc -o tools/ --gen-onefile --gen-object-api --python src/schema.fbs
+rten-convert/rten_convert/schema_generated.py: src/schema.fbs
+	flatc -o rten-convert/rten_convert --gen-onefile --gen-object-api --python src/schema.fbs
 
 
 .PHONY: gen-pytorch-references

README.md

@@ -50,14 +50,17 @@ classification example:
 git clone https://github.com/robertknight/rten.git
 cd rten
 
+# Install dependencies for model conversion
+pip install -e rten-convert
+
 # Install dependencies for Python scripts
 pip install -r tools/requirements.txt
 
 # Export an ONNX model. We're using resnet-50, a classic image classification model.
 python -m tools.export-timm-model timm/resnet50.a1_in1k
 
 # Convert model to this library's format
-tools/convert-onnx.py resnet50.a1_in1k.onnx resnet50.rten
+rten-convert resnet50.a1_in1k.onnx resnet50.rten
 
 # Run image classification example. Replace `image.png` with your own image.
 cargo run -p rten-examples --release --bin imagenet mobilenet resnet50.rten image.png
@@ -76,8 +79,8 @@ run:
 
 ```sh
 git clone https://github.com/robertknight/rten.git
-pip install -r rten/tools/requirements.txt
-rten/tools/convert-onnx.py your-model.onnx output.rten
+pip install -e rten/rten-convert
+rten-convert your-model.onnx output.rten
 ```
 
 The RTen model format does not yet guarantee long-term backwards compatibility,
@@ -111,7 +114,7 @@ a JavaScript project are:
     or [Hugging Face](https://huggingface.co/docs/transformers/serialization).
 2.  If the model is not already in ONNX format, convert it to ONNX. PyTorch
     users can use [torch.onnx](https://pytorch.org/docs/stable/onnx.html) for this.
-3.  Use the `tools/convert-onnx.py` script in this repository to convert the model
+3.  Use the `rten-convert` package in this repository to convert the model
     to the optimized format RTen uses. See the section above on converting models.
 4.  In your JavaScript code, fetch the WebAssembly binary and initialize RTen
     using the `init` function.

docs/adding-operators.md

@@ -19,7 +19,7 @@ In detail, the process is:
    attributes from that operator.
 3. Run `make` to generate updated Rust and Python code to read the updated
    FlatBuffers schema
-4. If the new operator has attributes, edit `tools/convert-onnx.py` to read
+4. If the new operator has attributes, edit `rten-convert/rten_convert/converter.py` and reinstall rten-convert to read
    the attributes from ONNX and convert to this library's model format
 5. Define the implementation of the new operator in Rust. This is a struct
    that implements the `Operator` trait.

js-examples/image-classification/README.md

@@ -14,7 +14,8 @@ kind of object.
 
    ```sh
    curl -L https://github.com/onnx/models/raw/main/vision/classification/mobilenet/model/mobilenetv2-10.onnx -o mobilenet.onnx
-   ../../tools/convert-onnx.py mobilenet.onnx mobilenet.rten
+
+   rten-convert mobilenet.onnx mobilenet.rten
    ```
 4. Follow either of the subsections below to run the example in Node or the
    browser

rten-cli/src/main.rs

@@ -293,7 +293,7 @@ fn print_input_output_list(model: &Model, node_ids: &[NodeId]) {
 /// generated inputs.
 ///
 /// ```
-/// tools/convert-onnx.py model.onnx output.rten
+/// rten-convert model.onnx output.rten
 /// cargo run -p rten-cli --release output.rten
 /// ```
 ///

rten-convert/README.md

@@ -0,0 +1,7 @@
+The conversion tool requires Python >= 3.10. To convert an existing ONNX model,
+run:
+
+```sh
+pip install -e .
+rten-convert your-model.onnx output.rten
+```

rten-convert/pyproject.toml

@@ -0,0 +1,8 @@
+[project]
+name = "rten-convert"
+requires-python = ">=3.10"
+version = "0.4.0"
+dependencies = ["flatbuffers", "onnx", "numpy", "setuptools"]
+
+[project.scripts]
+rten-convert = "rten_convert.converter:main"

tools/convert-onnx.py → rten-convert/rten_convert/converter.py

@@ -14,7 +14,7 @@
 import onnx.numpy_helper as numpy_helper
 from onnx import TensorProto, ValueInfoProto
 
-import schema_generated as sg
+import rten_convert.schema_generated as sg
 
 AttributeValue = int | float | str | list[int]
 

tools/schema_generated.py → rten-convert/rten_convert/schema_generated.py

@@ -1,6 +1,6 @@
 # automatically generated by the FlatBuffers compiler, do not modify
 
-# namespace: 
+# namespace:
 
 import flatbuffers
 from flatbuffers.compat import import_numpy
@@ -5066,5 +5066,3 @@ def Pack(self, builder):
             ModelAddMetadata(builder, metadata)
         model = ModelEnd(builder)
         return model
-
-

rten-examples/README.md

@@ -16,11 +16,11 @@ The general steps to run an example are:
    Face](https://huggingface.co/docs/optimum/exporters/onnx/overview),
    the [ONNX Model Zoo](https://github.com/onnx/models) or pre-created ONNX
    models by the model authors.
-2. Convert the ONNX model to this library's format using the `convert-onnx.py`
-   script:
+2. Convert the ONNX model to this library's format using the `rten-convert`
+   package:
 
    ```sh
-   $ ../tools/convert-onnx.py <onnx_model> <output_model>
+   $ rten-convert <onnx_model> <output_model>
    ```
 
 3. Run the example using:

rten-examples/src/bert_qa.rs

@@ -206,7 +206,7 @@ fn extract_nbest_answers<'a>(
 ///
 /// ```
 /// optimum-cli export onnx --model distilbert-base-cased-distilled-squad distilbert
-/// tools/convert-onnx.py distilbert/model.onnx distilbert/distilbert.rten
+/// rten-convert distilbert/model.onnx distilbert/distilbert.rten
 /// ```
 ///
 /// Then run the example with:

rten-examples/src/deeplab.rs

@@ -95,7 +95,7 @@ const PASCAL_VOC_LABELS: [(&str, Rgb); 21] = [
 ///
 /// ```
 /// python examples/export-deeplab.py
-/// tools/convert-onnx.py deeplab.onnx deeplab.rten
+/// rten-convert deeplab.onnx deeplab.rten
 /// ```
 ///
 /// Run this program on an image with:

rten-examples/src/depth_anything.rs

@@ -66,7 +66,7 @@ Args:
 /// After downloading the model, it can be run on an image using:
 ///
 /// ```
-/// tools/convert-onnx.py depth_anything.onnx
+/// rten-convert depth_anything.onnx
 /// cargo run --release --bin depth_anything depth_anything.rten image.jpg
 /// ```
 ///

rten-examples/src/detr.rs

@@ -261,7 +261,7 @@ const LABELS: &[&str] = &[
 ///
 /// ```
 /// optimum-cli export onnx --model facebook/detr-resnet-50 detr
-/// tools/convert-onnx.py detr/model.onnx detr.rten
+/// rten-convert detr/model.onnx detr.rten
 /// ```
 ///
 /// This model also works with YOLOS. Use `hustvl/yolos-tiny` or

rten-examples/src/jina_similarity.rs

@@ -176,7 +176,7 @@ fn embed_sentence_batch(
 /// Convert the model using:
 ///
 /// ```text
-/// tools/convert-onnx.py jina-embed.onnx jina-embed.rten
+/// rten-convert jina-embed.onnx jina-embed.rten
 /// ```
 ///
 /// Then run the example with:

rten-examples/src/wav2vec2.rs

@@ -78,7 +78,7 @@ fn read_wav_file(path: &str) -> Result<Vec<f32>, hound::Error> {
 ///
 /// ```
 /// optimum-cli export onnx --model facebook/wav2vec2-base-960h wav2vec2
-/// tools/convert-onnx.py wav2vec2/model.onnx wav2vec2.rten
+/// rten-convert wav2vec2/model.onnx wav2vec2.rten
 /// ```
 ///
 /// To record a .wav file and test this app:

rten-examples/src/yolo.rs

@@ -80,7 +80,7 @@ fn resource_path(path: &str) -> PathBuf {
 /// ```
 /// pip install ultralytics
 /// yolo mode=export model=yolov8s.pt format=onnx
-/// tools/convert-onnx.py yolov8s.onnx yolov8.rten
+/// rten-convert yolov8s.onnx yolov8.rten
 /// ```
 ///
 /// Run this program on an image: