From fab99fa3b9f8e3b3beac16a89d1e35739501bfeb Mon Sep 17 00:00:00 2001 From: Torsten Woertwein Date: Thu, 2 Jan 2025 09:48:06 -0500 Subject: [PATCH 1/4] Example to export model to ONNX --- examples/deploy_as_onnx.py | 157 +++++++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 examples/deploy_as_onnx.py diff --git a/examples/deploy_as_onnx.py b/examples/deploy_as_onnx.py new file mode 100644 index 00000000..014a4893 --- /dev/null +++ b/examples/deploy_as_onnx.py @@ -0,0 +1,157 @@ +"""Convert a simple RSMTool model to ONNX.""" + +from pathlib import Path + +import numpy as np +from onnxruntime import InferenceSession + + +def convert( + model_file: Path, + feature_file: Path, + calibration_file: Path, + trim_min: float, + trim_max: float, + trim_tolerance: float, + verify_correctness: bool = True, +) -> None: + # Convert a simple rsmtool model to onnx. + import json + + import pandas as pd + from skl2onnx import to_onnx + from skll.learner import Learner + + # load files + learner = Learner.from_file(model_file) + feature_information = pd.read_csv(feature_file, index_col=0) + calibrated_values = json.loads(Path(calibration_file).read_text()) + + # validate simplifying assumptions + assert (feature_information["transform"] == "raw").all(), "Only transform=raw is implemented" + assert (feature_information["sign"] == 1).all(), "Only sign=1 is implemented" + assert learner.feat_selector.get_support().all(), "Remove features from df_feature_info" + + # sort features names (FeatureSet does that) + feature_information = feature_information.sort_values(by="feature") + + # combine calibration values into one transformation + scale = calibrated_values["human_labels_sd"] / calibrated_values["train_predictions_sd"] + shift = ( + calibrated_values["human_labels_mean"] - calibrated_values["train_predictions_mean"] * scale + ) + + # export model and statistics + onnx_model = to_onnx( + learner.model, + feature_information["train_mean"].to_numpy().astype(np.float32)[None], + target_opset=20, + ) + model_file.with_suffix(".onnx").write_bytes(onnx_model.SerializeToString()) + + statistics = { + "feature_names": feature_information.index.to_list(), + "feature_outlier_min": ( + feature_information["train_mean"] - 4 * feature_information["train_sd"] + ).to_list(), + "feature_outlier_max": ( + feature_information["train_mean"] + 4 * feature_information["train_sd"] + ).to_list(), + "feature_means": feature_information["train_transformed_mean"].to_list(), + "feature_stds": feature_information["train_transformed_sd"].to_list(), + "label_mean": shift, + "label_std": scale, + "label_min": trim_min - trim_tolerance, + "label_max": trim_max + trim_tolerance, + } + (model_file.parent / f"{model_file.with_suffix('').name}_statistics.json").write_text( + json.dumps(statistics) + ) + + if not verify_correctness: + return + + # verify that the converted model produces the same output + from time import time + + from rsmtool import fast_predict + from rsmtool.modeler import Modeler + + onnx_model = InferenceSession(model_file.with_suffix(".onnx")) + rsm_model = Modeler.load_from_learner(learner) + onnx_duration = 0 + rsm_duration = 0 + iterations = 1_000 + for _ in range(iterations): + # sample random input data + features = ( + feature_information["train_mean"] + + (np.random.rand(feature_information.shape[0]) - 0.5) + * 10 + * feature_information["train_sd"] + ).to_dict() + + start = time() + onnx_prediction = predict(features, model=onnx_model, statistics=statistics) + onnx_duration += time() - start + + start = time() + rsm_prediction = fast_predict( + features, + modeler=rsm_model, + df_feature_info=feature_information, + trim=True, + trim_min=trim_min, + trim_max=trim_max, + trim_tolerance=trim_tolerance, + scale=True, + train_predictions_mean=calibrated_values["train_predictions_mean"], + train_predictions_sd=calibrated_values["train_predictions_sd"], + h1_mean=calibrated_values["human_labels_mean"], + h1_sd=calibrated_values["human_labels_sd"], + )["scale_trim"] + rsm_duration += time() - start + + assert np.isclose(onnx_prediction, rsm_prediction) + + print(f"ONNX duration: {round(onnx_duration/iterations, 5)}") + print(f"RSMTool duration: {round(rsm_duration/iterations, 5)}") + + +def predict( + features: dict[str, float], + *, + model: InferenceSession, + statistics: dict[str, np.ndarray | float | list[str]], +) -> float: + # get features in the expected order + features = np.array([features[name] for name in statistics["feature_names"]]) + + # clip outliers + features = np.clip( + features, a_min=statistics["feature_outlier_min"], a_max=statistics["feature_outlier_max"] + ) + + # normalize + features = (features - statistics["feature_means"]) / statistics["feature_stds"] + + # predict + # prediction = model.predict(features[None]) + prediction = model.run(None, {"X": features[None].astype(np.float32)})[0].item() + + # transform to human scale + prediction = prediction * statistics["label_std"] + statistics["label_mean"] + + # trim both raw and scaled predictions if requested + return np.clip(prediction, a_min=statistics["label_min"], a_max=statistics["label_max"]) + + +if __name__ == "__main__": + convert( + Path("test.model"), + Path("features.csv"), + Path("calibrated_values.json"), + trim_min=1, + trim_max=3, + trim_tolerance=0.49998, + ) From 1a6d7bb62c9ee4580e3777037d73dd3e341cd9e0 Mon Sep 17 00:00:00 2001 From: Torsten Woertwein Date: Thu, 2 Jan 2025 10:08:45 -0500 Subject: [PATCH 2/4] doc-strings --- examples/deploy_as_onnx.py | 37 +++++++++++++++++++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/examples/deploy_as_onnx.py b/examples/deploy_as_onnx.py index 014a4893..34592091 100644 --- a/examples/deploy_as_onnx.py +++ b/examples/deploy_as_onnx.py @@ -15,7 +15,26 @@ def convert( trim_tolerance: float, verify_correctness: bool = True, ) -> None: - # Convert a simple rsmtool model to onnx. + """Convert a simple rsmtool model to onnx. + + Parameters + ---------- + model_file: + Path to the file containing the SKLL learner. + feature_file: + Path to the file containing the feature statistics. + calibration_file: + Path to the file containing the label statistics. + trim_min,trim_max,trim_tolerance: + Trimming arguments for `fast_predict`. + verify_correctness: + Whether to verify that the converted model produces the same output. + + Raises + ------ + AssertionError + If an unsupported operation is encountered or the correctness test failed. + """ import json import pandas as pd @@ -124,6 +143,21 @@ def predict( model: InferenceSession, statistics: dict[str, np.ndarray | float | list[str]], ) -> float: + """Make a single prediction with the convered ONNX model. + + Parameters + ---------- + features: + Dictionary of the input features. + model: + ONNX inference session of the converted model. + statistics: + Dictionary containing the feature and label statistics. + + Returns + ------- + A single prediction. + """ # get features in the expected order features = np.array([features[name] for name in statistics["feature_names"]]) @@ -136,7 +170,6 @@ def predict( features = (features - statistics["feature_means"]) / statistics["feature_stds"] # predict - # prediction = model.predict(features[None]) prediction = model.run(None, {"X": features[None].astype(np.float32)})[0].item() # transform to human scale From 1cc60e03ed018989ce8d05d3ba16331800556cb9 Mon Sep 17 00:00:00 2001 From: Torsten Woertwein Date: Thu, 2 Jan 2025 11:57:15 -0500 Subject: [PATCH 3/4] usage example --- doc/advanced_usage.rst | 2 ++ doc/usage_onnx_deployment.rst.inc | 16 ++++++++++++++++ examples/deploy_as_onnx.py | 4 ++-- 3 files changed, 20 insertions(+), 2 deletions(-) create mode 100644 doc/usage_onnx_deployment.rst.inc diff --git a/doc/advanced_usage.rst b/doc/advanced_usage.rst index d7ababc8..5e8cc9bd 100644 --- a/doc/advanced_usage.rst +++ b/doc/advanced_usage.rst @@ -16,3 +16,5 @@ In addition to providing the ``rsmtool`` utility training and evaluating regress .. include:: usage_rsmxval.rst.inc .. include:: usage_rsmexplain.rst.inc + +.. include:: usage_onnx_deployment.rst.inc diff --git a/doc/usage_onnx_deployment.rst.inc b/doc/usage_onnx_deployment.rst.inc new file mode 100644 index 00000000..f2488f91 --- /dev/null +++ b/doc/usage_onnx_deployment.rst.inc @@ -0,0 +1,16 @@ +.. _usage_onnx_deployment: + +Deploy RSMTool models +^^^^^^^^^^^^^^^^^^^^^ + +RSMTool depends on many large python libraries which can make it tricky to efficiently deploy the trained models. This example `deploy_as_onnx.py `_ demonstrates how to export a simple RSMTool model to ONNX. The resulting model depends only on ``onnxruntime`` and ``numpy``. + +Pre- and post-processing +"""""""""""""""""""""""" + +The example script supports many pre-processing steps of RSMTool, such as, clipping outliers and z-normalization, and also post-processing steps, such as, scaling and clipping predictions. These steps are done in numpy, before and after calling the ONNX model. While not all features of RSMTool are supported, many of them could be supported by adjusting the numpy pre- and post-processing code. + +Model export +"""""""""""" + +In this example, we use `skl2onnx https://pypi.org/project/skl2onnx/`_ to export the underlying scikit-learn model to ONNX. Should this process fail, it is possible to export the scikit-learn model with ``joblib`` (``scikit-learn`` will then be a runtime dependecy). diff --git a/examples/deploy_as_onnx.py b/examples/deploy_as_onnx.py index 34592091..d68c021a 100644 --- a/examples/deploy_as_onnx.py +++ b/examples/deploy_as_onnx.py @@ -131,7 +131,7 @@ def convert( )["scale_trim"] rsm_duration += time() - start - assert np.isclose(onnx_prediction, rsm_prediction) + assert np.isclose(onnx_prediction, rsm_prediction), f"{onnx_prediction} vs {rsm_prediction}" print(f"ONNX duration: {round(onnx_duration/iterations, 5)}") print(f"RSMTool duration: {round(rsm_duration/iterations, 5)}") @@ -175,7 +175,7 @@ def predict( # transform to human scale prediction = prediction * statistics["label_std"] + statistics["label_mean"] - # trim both raw and scaled predictions if requested + # trim prediction return np.clip(prediction, a_min=statistics["label_min"], a_max=statistics["label_max"]) From f0a28cccf0287b1dbbe7705e3de0b51773a0912e Mon Sep 17 00:00:00 2001 From: Torsten Woertwein Date: Thu, 2 Jan 2025 12:00:49 -0500 Subject: [PATCH 4/4] comment on correctness --- doc/usage_onnx_deployment.rst.inc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/usage_onnx_deployment.rst.inc b/doc/usage_onnx_deployment.rst.inc index f2488f91..7f31b9c9 100644 --- a/doc/usage_onnx_deployment.rst.inc +++ b/doc/usage_onnx_deployment.rst.inc @@ -14,3 +14,7 @@ Model export """""""""""" In this example, we use `skl2onnx https://pypi.org/project/skl2onnx/`_ to export the underlying scikit-learn model to ONNX. Should this process fail, it is possible to export the scikit-learn model with ``joblib`` (``scikit-learn`` will then be a runtime dependecy). + +Correctness +""""""""""" +The example script calls the converted model with many different inputs to verify that it produces the same output as the original RSMTool model.