Skip to content

[onnx_utils] Upgrade versions and fixed tests (#853) #857

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Eyal-Danieli merged 1 commit into from
Apr 28, 2025
Merged

[onnx_utils] Upgrade versions and fixed tests (#853) #857

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 53 additions & 59 deletions onnx_utils/function.yaml
View file
Open in desktop

Large diffs are not rendered by default.

20 changes: 13 additions & 7 deletions onnx_utils/item.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ labels:
author: guyl
maintainers: []
marketplaceType: ''
mlrunVersion: 1.1.0
mlrunVersion: 1.7.2
name: onnx_utils
platformVersion: 3.5.0
spec:
Expand All @@ -26,10 +26,16 @@ spec:
image: mlrun/mlrun
kind: job
requirements:
- onnx~=1.13.0
- onnxruntime~=1.14.0
- onnxoptimizer~=0.3.0
- onnxmltools~=1.11.0
- tf2onnx~=1.13.0
- tqdm~=4.67.1
- tensorflow~=2.19.0
- tf_keras~=2.19.0
- torch~=2.6.0
- torchvision~=0.21.0
- onnx~=1.17.0
- onnxruntime~=1.19.2
- onnxoptimizer~=0.3.13
- onnxmltools~=1.13.0
- tf2onnx~=1.16.1
- plotly~=5.4.0
url: ''
version: 1.2.0
version: 1.3.0
67 changes: 34 additions & 33 deletions onnx_utils/onnx_utils.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,15 +84,14 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"TF_USE_LEGACY_KERAS\"] = \"true\"\n",
"from tempfile import TemporaryDirectory\n",
"\n",
"# Create a temporary directory for the model artifact:\n",
Expand All @@ -107,7 +106,9 @@
"\n",
"# Choose our optimized ONNX version model's name:\n",
"OPTIMIZED_ONNX_MODEL_NAME = \"optimized_onnx_mobilenetv2\""
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand All @@ -122,22 +123,20 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"# mlrun: start-code"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from tensorflow import keras\n",
"\n",
Expand All @@ -158,20 +157,22 @@
"\n",
" # Log the model:\n",
" model_handler.log()"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"# mlrun: end-code"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand All @@ -186,13 +187,11 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"import mlrun\n",
"\n",
Expand All @@ -213,7 +212,9 @@
" },\n",
" local=True\n",
")"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand All @@ -224,13 +225,11 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"# Import the ONNX function from the marketplace:\n",
"onnx_utils_function = mlrun.import_function(\"hub://onnx_utils\")\n",
Expand All @@ -247,7 +246,9 @@
" },\n",
" local=True\n",
")"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand All @@ -258,19 +259,19 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"import os\n",
"\n",
"\n",
"print(os.listdir(ARTIFACT_PATH))"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -304,13 +305,11 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"onnx_utils_function.run(\n",
" handler=\"optimize\",\n",
Expand All @@ -322,7 +321,9 @@
" },\n",
" local=True\n",
")"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand All @@ -333,16 +334,16 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"print(os.listdir(ARTIFACT_PATH))"
]
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
Expand All @@ -357,24 +358,24 @@
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"import shutil\n",
"\n",
"\n",
"shutil.rmtree(ARTIFACT_PATH)"
]
],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
Expand All @@ -388,9 +389,9 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.10"
"version": "3.9.22"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
}
32 changes: 20 additions & 12 deletions onnx_utils/onnx_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ def pytorch_to_onnx(
input_layers_names=input_layers_names,
output_layers_names=output_layers_names,
dynamic_axes=dynamic_axes,
is_batched=is_batched
is_batched=is_batched,
)


Expand All @@ -160,26 +160,31 @@ def pytorch_to_onnx(
def to_onnx(
context: mlrun.MLClientCtx,
model_path: str,
load_model_kwargs: dict = None,
onnx_model_name: str = None,
optimize_model: bool = True,
framework_kwargs: Dict[str, Any] = None,
):
"""
Convert the given model to an ONNX model.

:param context: The MLRun function execution context
:param model_path: The model path store object.
:param onnx_model_name: The name to use to log the converted ONNX model. If not given, the given `model_name` will
be used with an additional suffix `_onnx`. Defaulted to None.
:param optimize_model: Whether to optimize the ONNX model using 'onnxoptimizer' before saving the model. Defaulted
to True.
:param framework_kwargs: Additional arguments each framework may require in order to convert to ONNX. To get the doc
string of the desired framework onnx conversion function, pass "help".
:param context: The MLRun function execution context
:param model_path: The model path store object.
:param load_model_kwargs: Keyword arguments to pass to the `AutoMLRun.load_model` method.
:param onnx_model_name: The name to use to log the converted ONNX model. If not given, the given `model_name` will
be used with an additional suffix `_onnx`. Defaulted to None.
:param optimize_model: Whether to optimize the ONNX model using 'onnxoptimizer' before saving the model.
Defaulted to True.
:param framework_kwargs: Additional arguments each framework may require to convert to ONNX. To get the doc string
of the desired framework onnx conversion function, pass "help".
"""
from mlrun.frameworks.auto_mlrun.auto_mlrun import AutoMLRun

# Get a model handler of the required framework:
model_handler = AutoMLRun.load_model(model_path=model_path, context=context)
load_model_kwargs = load_model_kwargs or {}
model_handler = AutoMLRun.load_model(
model_path=model_path, context=context, **load_model_kwargs
)

# Get the model's framework:
framework = model_handler.FRAMEWORK_NAME
Expand Down Expand Up @@ -219,6 +224,7 @@ def to_onnx(
def optimize(
context: mlrun.MLClientCtx,
model_path: str,
handler_init_kwargs: dict = None,
optimizations: List[str] = None,
fixed_point: bool = False,
optimized_model_name: str = None,
Expand All @@ -228,8 +234,9 @@ def optimize(

:param context: The MLRun function execution context.
:param model_path: Path to the ONNX model object.
:param handler_init_kwargs: Keyword arguments to pass to the `ONNXModelHandler` init method preloading.
:param optimizations: List of possible optimizations. To see what optimizations are available, pass "help".
If None, all of the optimizations will be used. Defaulted to None.
If None, all the optimizations will be used. Defaulted to None.
:param fixed_point: Optimize the weights using fixed point. Defaulted to False.
:param optimized_model_name: The name of the optimized model. If None, the original model will be overridden.
Defaulted to None.
Expand All @@ -245,8 +252,9 @@ def optimize(
return

# Create the model handler:
handler_init_kwargs = handler_init_kwargs or {}
model_handler = ONNXModelHandler(
model_path=model_path, context=context
model_path=model_path, context=context, **handler_init_kwargs
)

# Load the ONNX model:
Expand Down
20 changes: 10 additions & 10 deletions onnx_utils/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
tqdm~=4.62.3
tensorflow~=2.7.0
torch~=1.10.0
torchvision~=0.11.1
onnx~=1.10.1
onnxruntime~=1.8.1
onnxoptimizer~=0.2.0
onnxmltools~=1.9.0
tf2onnx~=1.9.0
tqdm~=4.67.1
tensorflow~=2.19.0
tf_keras~=2.19.0
torch~=2.6.0
torchvision~=0.21.0
onnx~=1.17.0
onnxruntime~=1.19.2
onnxoptimizer~=0.3.13
onnxmltools~=1.13.0
tf2onnx~=1.16.1
plotly~=5.4.0
wrapt<1.15.0 # wrapt==1.15.0 fails tensorflow 2.7 Todo: please remove when updating tensorflow
Loading
Loading