Implment basic ONNX serialisation for models (#226)

Author: RAMitchell

conda/environments/all_cuda-122.yaml

@@ -28,6 +28,8 @@ dependencies:
 - ninja>=1.11.1.1
 - notebook>=7
 - numpy
+- onnx>=1.10
+- onnxmltools>=1.10
 - openblas
 - pydata-sphinx-theme>=0.16
 - pytest>=7,<8

dependencies.yaml

@@ -176,3 +176,5 @@ dependencies:
           - pytest>=7,<8
           - seaborn>=0.13
           - xgboost>=2.0
+          - onnx>=1.10
+          - onnxmltools>=1.10

legateboost/models/base_model.py

@@ -126,3 +126,13 @@ def __mul__(self, scalar: Any) -> "BaseModel":
 
     def __hash__(self) -> int:
         return hash(str(self))
+
+    def to_onnx(self) -> Any:
+        """Convert the model to an ONNX model.
+
+        Returns
+        -------
+        Any
+            The ONNX model.
+        """
+        raise NotImplementedError

legateboost/models/krr.py

@@ -242,3 +242,116 @@ def __mul__(self, scalar: Any) -> "KRR":
         new = copy.deepcopy(self)
         self.betas_ *= scalar
         return new
+
+    def to_onnx(self) -> Any:
+        from onnx import numpy_helper
+        from onnx.checker import check_model
+        from onnx.helper import (
+            make_graph,
+            make_model,
+            make_node,
+            make_tensor_value_info,
+            np_dtype_to_tensor_dtype,
+        )
+
+        assert self.X_train.dtype == self.betas_.dtype
+
+        def make_constant_node(value: cn.array, name: str) -> Any:
+            return make_node(
+                "Constant",
+                inputs=[],
+                value=numpy_helper.from_array(value, name=name),
+                outputs=[name],
+            )
+
+        nodes = []
+
+        # model constants
+        betas = numpy_helper.from_array(self.betas_.__array__(), name="betas")
+        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]],
+        )
+        pred = make_tensor_value_info(
+            "pred",
+            np_dtype_to_tensor_dtype(self.betas_.dtype),
+            [None, self.betas_.shape[1]],
+        )
+
+        # 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_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(
+            make_constant_node(
+                np.array([-2.0 * self.sigma**2], self.betas_.dtype), "denominator"
+            )
+        )
+        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"]))
+        graph = make_graph(
+            nodes, "legateboost.model.KRR", [X], [pred], [betas, X_train]
+        )
+        onnx_model = make_model(graph)
+        check_model(onnx_model)
+        return onnx_model

legateboost/models/linear.py

@@ -58,7 +58,7 @@ def __init__(
             self.l2_regularization = alpha
 
     def _fit_solve(self, X: cn.ndarray, g: cn.ndarray, h: cn.ndarray) -> None:
-        self.betas_ = cn.zeros((X.shape[1] + 1, g.shape[1]))
+        self.betas_ = cn.zeros((X.shape[1] + 1, g.shape[1]), dtype=X.dtype)
         num_outputs = g.shape[1]
         for k in range(num_outputs):
             W = cn.sqrt(h[:, k])
@@ -135,12 +135,13 @@ def batch_predict(models: Sequence[BaseModel], X: cn.ndarray) -> cn.ndarray:
         # summing together the coeffiecients of each model then predicting
         # saves a lot of work
         betas = cn.sum([model.betas_ for model in models], axis=0)
-        return betas[0] + X.dot(betas[1:].astype(X.dtype))
+        betas = betas.astype(X.dtype)
+        return betas[0] + X.dot(betas[1:])
 
     def __str__(self) -> str:
         return (
             "Bias: "
-            + str(self.betas_[1])
+            + str(self.betas_[0])
             + "\nCoefficients: "
             + str(self.betas_[1:])
             + "\n"
@@ -150,3 +151,37 @@ def __mul__(self, scalar: Any) -> "Linear":
         new = copy.deepcopy(self)
         new.betas_ *= scalar
         return new
+
+    def to_onnx(self) -> Any:
+        from onnx import numpy_helper
+        from onnx.checker import check_model
+        from onnx.helper import (
+            make_graph,
+            make_model,
+            make_node,
+            make_tensor_value_info,
+            np_dtype_to_tensor_dtype,
+        )
+
+        # model constants
+        betas = numpy_helper.from_array(self.betas_[1:].__array__(), name="betas")
+        intercept = numpy_helper.from_array(
+            self.betas_[0].__array__(), name="intercept"
+        )
+
+        # pred inputs
+        X = make_tensor_value_info(
+            "X", np_dtype_to_tensor_dtype(self.betas_.dtype), [None, None]
+        )
+        pred = make_tensor_value_info(
+            "pred", np_dtype_to_tensor_dtype(self.betas_.dtype), [None]
+        )
+
+        node1 = make_node("MatMul", ["X", "betas"], ["XBeta"])
+        node2 = make_node("Add", ["XBeta", "intercept"], ["pred"])
+        graph = make_graph(
+            [node1, node2], "legateboost.model.Linear", [X], [pred], [betas, intercept]
+        )
+        onnx_model = make_model(graph)
+        check_model(onnx_model)
+        return onnx_model

legateboost/models/nn.py

@@ -181,3 +181,87 @@ def __mul__(self, scalar: Any) -> "NN":
         new.coefficients_[-1] *= scalar
         new.biases_[-1] *= scalar
         return new
+
+    def to_onnx(self) -> Any:
+        from onnx import numpy_helper
+        from onnx.checker import check_model
+        from onnx.helper import (
+            make_graph,
+            make_model,
+            make_node,
+            make_tensor_value_info,
+            np_dtype_to_tensor_dtype,
+        )
+
+        # model constants
+        biases = [
+            numpy_helper.from_array(b[0].__array__(), name=f"bias{i}")
+            for i, b in enumerate(self.biases_)
+        ]
+        coefficients = [
+            numpy_helper.from_array(c.__array__(), name=f"coefficients{i}")
+            for i, c in enumerate(self.coefficients_)
+        ]
+
+        # pred inputs
+        X = make_tensor_value_info(
+            "X",
+            np_dtype_to_tensor_dtype(self.coefficients_[0].dtype),
+            [None, self.coefficients_[0].shape[0]],
+        )
+
+        nodes = []
+
+        make_tensor_value_info(
+            "activations0",
+            np_dtype_to_tensor_dtype(self.coefficients_[0].dtype),
+            [None, None],
+        )
+        nodes.append(make_node("MatMul", ["X", "coefficients0"], ["activations0"]))
+        activations_with_bias = make_tensor_value_info(
+            "activations0withbias",
+            np_dtype_to_tensor_dtype(self.coefficients_[0].dtype),
+            [None, None],
+        )
+        nodes.append(
+            make_node("Add", ["activations0", "bias0"], ["activations0withbias"])
+        )
+
+        for i in range(1, len(coefficients)):
+            make_tensor_value_info(
+                f"tanh{i}",
+                np_dtype_to_tensor_dtype(self.coefficients_[0].dtype),
+                [None, None],
+            )
+            nodes.append(make_node("Tanh", [f"activations{i-1}withbias"], [f"tanh{i}"]))
+            make_tensor_value_info(
+                f"activations{i}",
+                np_dtype_to_tensor_dtype(self.coefficients_[0].dtype),
+                [None, None],
+            )
+            nodes.append(
+                make_node(
+                    "MatMul", [f"tanh{i}", f"coefficients{i}"], [f"activations{i}"]
+                )
+            )
+            activations_with_bias = make_tensor_value_info(
+                f"activations{i}withbias",
+                np_dtype_to_tensor_dtype(self.coefficients_[0].dtype),
+                [None, None],
+            )
+            nodes.append(
+                make_node(
+                    "Add", [f"activations{i}", f"bias{i}"], [f"activations{i}withbias"]
+                )
+            )
+
+        graph = make_graph(
+            nodes,
+            "legateboost.model.NN",
+            [X],
+            [activations_with_bias],
+            biases + coefficients,
+        )
+        onnx_model = make_model(graph)
+        check_model(onnx_model)
+        return onnx_model