Implement ONNX model serialisation phase 2

conda/environments/all_cuda-122.yaml

@@ -30,6 +30,7 @@ dependencies:
 - numpy
 - onnx>=1.10
 - onnxmltools>=1.10
+- onnxruntime>=1.21
 - openblas
 - pydata-sphinx-theme>=0.16
 - pytest>=7,<8

dependencies.yaml

@@ -178,3 +178,4 @@ dependencies:
           - xgboost>=2.0
           - onnx>=1.10
           - onnxmltools>=1.10
+          - onnxruntime>=1.21

legateboost/legateboost.py

@@ -540,6 +540,102 @@ def dump_models(self) -> str:
             text += str(m)
         return text
 
+    def _make_onnx_init(self, X_dtype):
+        # turn self.model_init_ into an ONNX model
+        from onnx import numpy_helper
+        from onnx.checker import check_model
+        from onnx.helper import (
+            make_graph,
+            make_model,
+            make_node,
+            make_opsetid,
+            make_tensor_value_info,
+            np_dtype_to_tensor_dtype,
+        )
+
+        # model constants
+        X_in = make_tensor_value_info(
+            "X_in", np_dtype_to_tensor_dtype(X_dtype), [None, self.n_features_in_]
+        )
+        nodes = []
+        nodes.append(make_node("Shape", ["X_in"], ["n_rows"], end=1))
+        one = numpy_helper.from_array(np.array([1], dtype=np.int64), name="one")
+        nodes.append(make_node("Concat", ["n_rows", "one"], ["tile_repeat"], axis=0))
+        init = numpy_helper.from_array(
+            np.atleast_2d(self.model_init_.__array__().astype(X_dtype)), name="init"
+        )
+        prediction_out = make_tensor_value_info(
+            "predictions_out",
+            np_dtype_to_tensor_dtype(X_dtype),
+            [None, self.model_init_.shape[0]],
+        )
+        nodes.append(make_node("Tile", ["init", "tile_repeat"], ["predictions_out"]))
+        X_out = make_tensor_value_info(
+            "X_out",
+            np_dtype_to_tensor_dtype(X_dtype),
+            [None, self.model_init_.shape[0]],
+        )
+        nodes.append(make_node("Identity", ["X_in"], ["X_out"]))
+        graph = make_graph(
+            nodes,
+            "legateboost estimator init",
+            [X_in],
+            [X_out, prediction_out],
+            [init, one],
+        )
+        onnx_model = make_model(
+            graph,
+            opset_imports=[
+                make_opsetid("", 21),
+            ],
+        )
+        check_model(onnx_model)
+
+        return onnx_model
+
+    def to_onnx(self, X_dtype, predict_function="predict"):
+        """Converts the model to an ONNX model.
+
+        Parameters
+        ----------
+        X_dtype : numpy.dtype
+            The expected data type of the input data. ONNX models hard
+            code the data type of the input data and will crash if this is
+            not set correctly.
+            Can be np.float32 or np.float64.
+        predict_function : str
+            The serialised ONNX model can produce output equivalent to 'predict',
+            'predict_proba', or 'predict_raw'.
+            The default is "predict".
+        Returns
+        -------
+        Any
+            The ONNX model.
+        """
+        from onnx.compose import merge_models
+
+        model = self._make_onnx_init(X_dtype)
+        if self.models_ is not None and len(self.models_) > 0:
+            model = merge_models(
+                model,
+                self.models_[0].to_onnx(X_dtype),
+                io_map=[("X_out", "X_in"), ("predictions_out", "predictions_in")],
+                prefix2="model_0_",
+            )
+
+        for i in range(1, len(self.models_)):
+            model = merge_models(
+                model,
+                self.models_[i].to_onnx(X_dtype),
+                io_map=[
+                    ("model_{}_X_out".format(i - 1), "X_in"),
+                    ("model_{}_predictions_out".format(i - 1), "predictions_in"),
+                ],
+                prefix2="model_{}_".format(i),
+            )
+
+        return model
+
     def global_attributions(
         self,
         X: cn.array,

legateboost/models/krr.py

@@ -243,13 +243,14 @@ def __mul__(self, scalar: Any) -> "KRR":
         self.betas_ *= scalar
         return new
 
-    def to_onnx(self) -> Any:
+    def to_onnx(self, X_dtype) -> Any:
         from onnx import numpy_helper
         from onnx.checker import check_model
         from onnx.helper import (
             make_graph,
             make_model,
             make_node,
+            make_opsetid,
             make_tensor_value_info,
             np_dtype_to_tensor_dtype,
         )
@@ -271,66 +272,34 @@ def make_constant_node(value: cn.array, name: str) -> Any:
         X_train = numpy_helper.from_array(self.X_train.__array__(), name="X_train")
 
         # pred inputs
-        X = make_tensor_value_info(
-            "X",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.X_train.shape[1]],
+        n_features = self.X_train.shape[1]
+        n_outputs = self.betas_.shape[1]
+        X_in = make_tensor_value_info(
+            "X_in", np_dtype_to_tensor_dtype(self.betas_.dtype), [None, n_features]
         )
-        pred = make_tensor_value_info(
-            "pred",
+        predictions_in = make_tensor_value_info(
+            "predictions_in",
             np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.betas_.shape[1]],
+            [None, n_outputs],
         )
-
         # exanded l2 distance
         # distance = np.sum(X**2, axis=1)[:, np.newaxis] - 2 * np.dot(X, self.X_train.T)
         # + np.sum(self.X_train**2, axis=1)
-        make_tensor_value_info(
-            "XX", np_dtype_to_tensor_dtype(self.betas_.dtype), [None]
-        )
-        make_tensor_value_info(
-            "YY",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [self.X_train.shape[0], 1],
-        )
-        make_tensor_value_info(
-            "XY_reshaped",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [1, self.X_train.shape[0]],
-        )
-        make_tensor_value_info(
-            "XY",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.X_train.shape[0]],
-        )
         nodes.append(make_constant_node(np.array([1]), "axis1"))
-        nodes.append(make_node("ReduceSumSquare", ["X", "axis1"], ["XX"]))
-        nodes.append(make_node("Gemm", ["X", "X_train"], ["XY"], alpha=-2.0, transB=1))
+        nodes.append(make_node("ReduceSumSquare", ["X_in", "axis1"], ["XX"]))
+        nodes.append(
+            make_node("Gemm", ["X_in", "X_train"], ["XY"], alpha=-2.0, transB=1)
+        )
         nodes.append(make_node("ReduceSumSquare", ["X_train", "axis1"], ["YY"]))
         nodes.append(make_constant_node(np.array([1, -1]), "reshape"))
         nodes.append(make_node("Reshape", ["YY", "reshape"], ["YY_reshaped"]))
         nodes.append(make_node("Add", ["XX", "XY"], ["add0"]))
-        make_tensor_value_info(
-            "l2",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.X_train.shape[0]],
-        )
         nodes.append(make_node("Add", ["YY_reshaped", "add0"], ["l2"]))
         nodes.append(make_constant_node(np.array([0.0], self.betas_.dtype), "zero"))
-        make_tensor_value_info(
-            "l2_clipped",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.X_train.shape[0]],
-        )
         nodes.append(make_node("Max", ["l2", "zero"], ["l2_clipped"]))
 
         # RBF kernel
         # K = np.exp(-distance / (2 * self.sigma**2))
-        make_tensor_value_info(
-            "rbf0",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.X_train.shape[0]],
-        )
         if self.sigma is None:
             raise ValueError("sigma is None. Has fit been called?")
         nodes.append(
@@ -339,19 +308,37 @@ def make_constant_node(value: cn.array, name: str) -> Any:
             )
         )
         nodes.append(make_node("Div", ["l2_clipped", "denominator"], ["rbf0"]))
-        make_tensor_value_info(
-            "K",
-            np_dtype_to_tensor_dtype(self.betas_.dtype),
-            [None, self.X_train.shape[0]],
-        )
         nodes.append(make_node("Exp", ["rbf0"], ["K"]))
 
         # prediction
         # pred = np.dot(K, self.betas_)
-        nodes.append(make_node("MatMul", ["K", "betas"], ["pred"]))
+        nodes.append(make_node("MatMul", ["K", "betas"], ["dot"]))
+
+        # outputs
+        predictions_out = make_tensor_value_info(
+            "predictions_out",
+            np_dtype_to_tensor_dtype(self.betas_.dtype),
+            [None, n_outputs],
+        )
+        X_out = make_tensor_value_info(
+            "X_out", np_dtype_to_tensor_dtype(self.betas_.dtype), [None, n_features]
+        )
+
+        nodes.append(make_node("Add", ["dot", "predictions_in"], ["predictions_out"]))
+        nodes.append(make_node("Identity", ["X_in"], ["X_out"]))
+
         graph = make_graph(
-            nodes, "legateboost.model.KRR", [X], [pred], [betas, X_train]
+            nodes,
+            "legateboost.model.KRR",
+            [X_in, predictions_in],
+            [X_out, predictions_out],
+            [betas, X_train],
+        )
+        onnx_model = make_model(
+            graph,
+            opset_imports=[
+                make_opsetid("", 21),
+            ],
         )
-        onnx_model = make_model(graph)
         check_model(onnx_model)
         return onnx_model

legateboost/models/linear.py

@@ -152,13 +152,14 @@ def __mul__(self, scalar: Any) -> "Linear":
         new.betas_ *= scalar
         return new
 
-    def to_onnx(self) -> Any:
+    def to_onnx(self, X_dtype) -> Any:
         from onnx import numpy_helper
         from onnx.checker import check_model
         from onnx.helper import (
             make_graph,
             make_model,
             make_node,
+            make_opsetid,
             make_tensor_value_info,
             np_dtype_to_tensor_dtype,
         )
@@ -170,18 +171,44 @@ def to_onnx(self) -> Any:
         )
 
         # pred inputs
-        X = make_tensor_value_info(
-            "X", np_dtype_to_tensor_dtype(self.betas_.dtype), [None, None]
+        n_features = self.betas_.shape[0] - 1
+        n_outputs = self.betas_.shape[1]
+        X_in = make_tensor_value_info(
+            "X_in", np_dtype_to_tensor_dtype(self.betas_.dtype), [None, n_features]
         )
-        pred = make_tensor_value_info(
-            "pred", np_dtype_to_tensor_dtype(self.betas_.dtype), [None]
+        predictions_in = make_tensor_value_info(
+            "predictions_in",
+            np_dtype_to_tensor_dtype(self.betas_.dtype),
+            [None, n_outputs],
+        )
+        predictions_out = make_tensor_value_info(
+            "predictions_out",
+            np_dtype_to_tensor_dtype(self.betas_.dtype),
+            [None, n_outputs],
         )
 
-        node1 = make_node("MatMul", ["X", "betas"], ["XBeta"])
-        node2 = make_node("Add", ["XBeta", "intercept"], ["pred"])
+        nodes = []
+        nodes.append(make_node("MatMul", ["X_in", "betas"], ["XBeta"]))
+        nodes.append(make_node("Add", ["XBeta", "intercept"], ["result"]))
+        nodes.append(
+            make_node("Add", ["result", "predictions_in"], ["predictions_out"])
+        )
+        X_out = make_tensor_value_info(
+            "X_out", np_dtype_to_tensor_dtype(self.betas_.dtype), [None, n_features]
+        )
+        nodes.append(make_node("Identity", ["X_in"], ["X_out"]))
         graph = make_graph(
-            [node1, node2], "legateboost.model.Linear", [X], [pred], [betas, intercept]
+            nodes,
+            "legateboost.model.Linear",
+            [X_in, predictions_in],
+            [X_out, predictions_out],
+            [betas, intercept],
+        )
+        onnx_model = make_model(
+            graph,
+            opset_imports=[
+                make_opsetid("", 21),
+            ],
         )
-        onnx_model = make_model(graph)
         check_model(onnx_model)
         return onnx_model