diff --git a/.gitignore b/.gitignore index f6f72add..5b0e7fc2 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ __pycache__ /dist /node_modules /target +/rten-convert/rten_convert.egg-info # Optimized neural network models *.model diff --git a/Makefile b/Makefile index 38223a5f..59c1b26e 100644 --- a/Makefile +++ b/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 diff --git a/README.md b/README.md index 6541980d..484764b7 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,9 @@ 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 @@ -57,7 +60,7 @@ pip install -r tools/requirements.txt 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. diff --git a/docs/adding-operators.md b/docs/adding-operators.md index a4a2950c..220fe5ea 100644 --- a/docs/adding-operators.md +++ b/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. diff --git a/js-examples/image-classification/README.md b/js-examples/image-classification/README.md index eacdfb5b..56c65cd5 100644 --- a/js-examples/image-classification/README.md +++ b/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 diff --git a/rten-cli/src/main.rs b/rten-cli/src/main.rs index cb02048d..2f6aba77 100644 --- a/rten-cli/src/main.rs +++ b/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 /// ``` /// diff --git a/rten-convert/README.md b/rten-convert/README.md new file mode 100644 index 00000000..f0a81207 --- /dev/null +++ b/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 +``` diff --git a/rten-convert/pyproject.toml b/rten-convert/pyproject.toml new file mode 100644 index 00000000..64c1b950 --- /dev/null +++ b/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" diff --git a/rten-convert/rten_convert/__init__.py b/rten-convert/rten_convert/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tools/convert-onnx.py b/rten-convert/rten_convert/converter.py old mode 100755 new mode 100644 similarity index 99% rename from tools/convert-onnx.py rename to rten-convert/rten_convert/converter.py index 2cadaf73..499e76a1 --- a/tools/convert-onnx.py +++ b/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] diff --git a/tools/schema_generated.py b/rten-convert/rten_convert/schema_generated.py similarity index 99% rename from tools/schema_generated.py rename to rten-convert/rten_convert/schema_generated.py index 0c74d47f..24b54414 100644 --- a/tools/schema_generated.py +++ b/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 - - diff --git a/rten-examples/README.md b/rten-examples/README.md index 0ab8a333..0e706f87 100644 --- a/rten-examples/README.md +++ b/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 + $ rten-convert ``` 3. Run the example using: diff --git a/rten-examples/src/bert_qa.rs b/rten-examples/src/bert_qa.rs index de523f35..f432a985 100644 --- a/rten-examples/src/bert_qa.rs +++ b/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: diff --git a/rten-examples/src/deeplab.rs b/rten-examples/src/deeplab.rs index 744ef5b1..94915029 100644 --- a/rten-examples/src/deeplab.rs +++ b/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: diff --git a/rten-examples/src/depth_anything.rs b/rten-examples/src/depth_anything.rs index 95bda4f6..5074e446 100644 --- a/rten-examples/src/depth_anything.rs +++ b/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 /// ``` /// diff --git a/rten-examples/src/detr.rs b/rten-examples/src/detr.rs index 81548077..ef4f8ac9 100644 --- a/rten-examples/src/detr.rs +++ b/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 diff --git a/rten-examples/src/jina_similarity.rs b/rten-examples/src/jina_similarity.rs index 9c9ccd50..e77b5228 100644 --- a/rten-examples/src/jina_similarity.rs +++ b/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: diff --git a/rten-examples/src/wav2vec2.rs b/rten-examples/src/wav2vec2.rs index f7f4313d..c320678f 100644 --- a/rten-examples/src/wav2vec2.rs +++ b/rten-examples/src/wav2vec2.rs @@ -78,7 +78,7 @@ fn read_wav_file(path: &str) -> Result, 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: diff --git a/rten-examples/src/yolo.rs b/rten-examples/src/yolo.rs index 27a768de..42d033ea 100644 --- a/rten-examples/src/yolo.rs +++ b/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: