kinfer 0.3.3__cp312-cp312-macosx_11_0_arm64.whl → 0.4.0__cp312-cp312-macosx_11_0_arm64.whl

kinfer/__init__.py

@@ -1,10 +1,5 @@
 """Defines the kinfer API."""
 
-from . import proto as K
-from .export.pytorch import export_model, get_model
-from .inference.base import KModel
-from .inference.python import ONNXModel
 from .rust_bindings import get_version
-from .serialize import get_multi_serializer, get_serializer
 
 __version__ = get_version()

kinfer/common/types.py

@@ -0,0 +1,11 @@
+"""Defines common types."""
+
+__all__ = [
+    "Metadata",
+]
+
+from pydantic import BaseModel
+
+
+class Metadata(BaseModel):
+    joint_names: list[str]

kinfer/export/common.py

@@ -0,0 +1,35 @@
+"""Defines common utilities for exporting models."""
+
+
+def get_shape(
+    name: str,
+    num_joints: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+) -> tuple[int, ...]:
+    match name:
+        case "joint_angles":
+            if num_joints is None:
+                raise ValueError("`num_joints` must be provided when using `joint_angles`")
+            return (num_joints,)
+
+        case "joint_angular_velocities":
+            if num_joints is None:
+                raise ValueError("`num_joints` must be provided when using `joint_angular_velocities`")
+            return (num_joints,)
+
+        case "projected_gravity":
+            return (3,)
+
+        case "accelerometer":
+            return (3,)
+
+        case "gyroscope":
+            return (3,)
+
+        case "carry":
+            if carry_shape is None:
+                raise ValueError("`carry_shape` must be provided for `carry`")
+            return carry_shape
+
+        case _:
+            raise ValueError(f"Unknown tensor name: {name}")

kinfer/export/jax.py

@@ -0,0 +1,51 @@
+"""Jax model export utilities."""
+
+import inspect
+import logging
+
+import tensorflow as tf
+import tf2onnx
+from equinox.internal._finalise_jaxpr import finalise_fn
+from jax._src.stages import Wrapped
+from jax.experimental import jax2tf
+from onnx.onnx_pb import ModelProto
+
+from kinfer.export.common import get_shape
+
+logger = logging.getLogger(__name__)
+
+
+def export_fn(
+    model: Wrapped,
+    *,
+    num_joints: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+    opset: int = 13,
+) -> ModelProto:
+    """Export a JAX function to ONNX."""
+    if not isinstance(model, Wrapped):
+        raise ValueError("Model must be a Wrapped function")
+
+    params = inspect.signature(model).parameters
+    input_names = list(params.keys())
+
+    # Gets the dummy input tensors for exporting the model.
+    tf_args = []
+    for name in input_names:
+        shape = get_shape(
+            name,
+            num_joints=num_joints,
+            carry_shape=carry_shape,
+        )
+        tf_args.append(tf.TensorSpec(shape, tf.float32, name=name))
+
+    finalised_fn = finalise_fn(model)
+    tf_fn = tf.function(jax2tf.convert(finalised_fn, enable_xla=False))
+
+    model_proto, _ = tf2onnx.convert.from_function(
+        tf_fn,
+        input_signature=tf_args,
+        opset=opset,
+        large_model=False,
+    )
+    return model_proto

kinfer/export/pytorch.py

@@ -1,128 +1,60 @@
 """PyTorch model export utilities."""
 
-import base64
-import inspect
-from io import BytesIO
-from typing import Sequence
+__all__ = [
+    "export_fn",
+]
+
+import io
+from typing import cast
 
 import onnx
-import onnxruntime as ort
 import torch
-from torch import Tensor
-
-from kinfer import proto as K
-from kinfer.serialize.pytorch import PyTorchMultiSerializer
-from kinfer.serialize.schema import get_dummy_io
-from kinfer.serialize.utils import check_names_match
+from onnx.onnx_pb import ModelProto
+from torch._C import FunctionSchema
 
-KINFER_METADATA_KEY = "kinfer_metadata"
+from kinfer.export.common import get_shape
 
 
-def _add_metadata_to_onnx(model_proto: onnx.ModelProto, schema: K.ModelSchema) -> onnx.ModelProto:
-    """Add metadata to ONNX model.
+def export_fn(
+    model: torch.jit.ScriptFunction,
+    *,
+    num_joints: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+) -> ModelProto:
+    """Exports a PyTorch function to ONNX.
 
     Args:
-        model_proto: ONNX model prototype
-        schema: Model schema to use for model export.
+        model: The model to export.
+        num_joints: The number of joints in the model.
+        carry_shape: The shape of the carry tensor.
 
     Returns:
-        ONNX model with added metadata
+        The ONNX model as a `ModelProto`.
     """
-    schema_bytes = schema.SerializeToString()
-
-    meta = model_proto.metadata_props.add()
-    meta.key = KINFER_METADATA_KEY
-    meta.value = base64.b64encode(schema_bytes).decode("utf-8")
-
-    return model_proto
-
-
-def export_model(model: torch.jit.ScriptModule, schema: K.ModelSchema) -> onnx.ModelProto:
-    """Export PyTorch model to ONNX format with metadata.
-
-    Args:
-        model: PyTorch model to export.
-        schema: Model schema to use for model export.
-
-    Returns:
-        ONNX inference session
-    """
-    # Matches each input name to the input values.
-    signature = inspect.signature(model.forward)
-    model_input_names = [
-        p.name for p in signature.parameters.values() if p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
-    ]
-    if len(model_input_names) != len(schema.input_schema.values):
-        raise ValueError(f"Expected {len(model_input_names)} inputs, but schema has {len(schema.input_schema.values)}")
-    input_schema_names = [i.value_name for i in schema.input_schema.values]
-    output_schema_names = [o.value_name for o in schema.output_schema.values]
-
-    if model_input_names != input_schema_names:
-        raise ValueError(f"Expected input names {model_input_names} to match schema names {input_schema_names}")
-
-    input_serializer = PyTorchMultiSerializer(schema.input_schema)
-    output_serializer = PyTorchMultiSerializer(schema.output_schema)
-
-    input_dummy_values = get_dummy_io(schema.input_schema)
-    input_tensors = input_serializer.serialize_io(input_dummy_values, as_dict=True)
-
-    check_names_match("model_input_names", model_input_names, "input_schema", list(input_tensors.keys()))
-    input_tensor_list = [input_tensors[name] for name in model_input_names]
-
-    # Attempts to run the model with the dummy inputs.
-    try:
-        pred_output_tensors = model(*input_tensor_list)
-    except Exception as e:
-        signature = inspect.signature(model.forward)
-        model_input_names = [
-            p.name for p in signature.parameters.values() if p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
-        ]
-
-        raise ValueError(
-            f"Failed to run model with dummy inputs; input names are {model_input_names} while "
-            f"input schema is {schema.input_schema}"
-        ) from e
-
-    # Attempts to parse the output tensors using the output schema.
-    if isinstance(pred_output_tensors, Tensor):
-        pred_output_tensors = (pred_output_tensors,)
-    if isinstance(pred_output_tensors, Sequence):
-        pred_output_tensors = output_serializer.assign_names(pred_output_tensors)
-    if not isinstance(pred_output_tensors, dict):
-        raise ValueError("Output tensors could not be converted to dictionary")
-    try:
-        pred_output_tensors = output_serializer.deserialize_io(pred_output_tensors)
-    except Exception as e:
-        raise ValueError("Failed to parse output tensors using output schema; are you sure it is correct?") from e
-
-    # Export model to buffer
-    buffer = BytesIO()
+    if not isinstance(model, torch.jit.ScriptFunction):
+        raise ValueError("Model must be a torch.jit.ScriptFunction")
+
+    schema = cast(FunctionSchema, model.schema)
+    input_names = [arg.name for arg in schema.arguments]
+
+    # Gets the dummy input tensors for exporting the model.
+    args = []
+    for name in input_names:
+        shape = get_shape(
+            name,
+            num_joints=num_joints,
+            carry_shape=carry_shape,
+        )
+        args.append(torch.zeros(shape))
+
+    buffer = io.BytesIO()
     torch.onnx.export(
         model=model,
         f=buffer,  # type: ignore[arg-type]
-        kwargs=input_tensors,
-        input_names=input_schema_names,
-        output_names=output_schema_names,
+        args=tuple(args),
+        input_names=input_names,
+        external_data=False,
     )
     buffer.seek(0)
-
-    # Loads the model from the buffer and adds metadata.
-    model_proto = onnx.load_model(buffer)
-    model_proto = _add_metadata_to_onnx(model_proto, schema)
-
-    return model_proto
-
-
-def get_model(model_proto: onnx.ModelProto) -> ort.InferenceSession:
-    """Converts a model proto to an inference session.
-
-    Args:
-        model_proto: ONNX model proto to convert to inference session.
-
-    Returns:
-        ONNX inference session
-    """
-    buffer = BytesIO()
-    onnx.save_model(model_proto, buffer)
-    buffer.seek(0)
-    return ort.InferenceSession(buffer.read())
+    model_bytes = buffer.read()
+    return onnx.load_from_string(model_bytes)

kinfer/export/serialize.py

@@ -0,0 +1,86 @@
+"""Functions for serializing and deserializing models."""
+
+__all__ = [
+    "pack",
+]
+
+
+import io
+import tarfile
+
+from onnx.onnx_pb import ModelProto
+
+from kinfer.common.types import Metadata
+from kinfer.export.common import get_shape
+
+
+def pack(
+    init_fn: ModelProto,
+    step_fn: ModelProto,
+    joint_names: list[str],
+    carry_shape: tuple[int, ...],
+) -> bytes:
+    """Packs the initialization function and step function into a directory.
+
+    Args:
+        init_fn: The initialization function.
+        step_fn: The step function.
+        joint_names: The list of joint names, in the order that the model
+            expects them to be provided.
+        carry_shape: The shape of the carry tensor.
+        root_dir: The root directory of the model.
+    """
+    num_joints = len(joint_names)
+
+    # Checks the `init` function.
+    if len(init_fn.graph.input) > 0:
+        raise ValueError(f"`init` function should not have any inputs! Got {len(init_fn.graph.input)}")
+    if len(init_fn.graph.output) != 1:
+        raise ValueError(f"`init` function should have exactly 1 output! Got {len(init_fn.graph.output)}")
+
+    # Checks the `step` function.
+    for step_input in step_fn.graph.input:
+        step_input_type = step_input.type.tensor_type
+        shape = tuple(dim.dim_value for dim in step_input_type.shape.dim)
+        expected_shape = get_shape(
+            step_input.name,
+            num_joints=num_joints,
+            carry_shape=carry_shape,
+        )
+        if shape != expected_shape:
+            raise ValueError(f"Expected shape {expected_shape} for input `{step_input.name}`, got {shape}")
+
+    if len(step_fn.graph.output) != 2:
+        raise ValueError(f"Step function must have exactly 2 outputs, got {len(step_fn.graph.output)}")
+
+    model_output = step_fn.graph.output[0]
+    output_shape = tuple(dim.dim_value for dim in model_output.type.tensor_type.shape.dim)
+    if output_shape != (num_joints,):
+        raise ValueError(f"Expected output shape {num_joints} for output `{model_output.name}`, got {output_shape}")
+
+    model_carry = step_fn.graph.output[1]
+    output_carry_shape = tuple(dim.dim_value for dim in model_carry.type.tensor_type.shape.dim)
+    if output_carry_shape != carry_shape:
+        raise ValueError(f"Expected carry shape {carry_shape} for output `{model_carry.name}`, got {carry_shape}")
+
+    # Builds the metadata object.
+    metadata = Metadata(
+        joint_names=joint_names,
+    )
+
+    buffer = io.BytesIO()
+
+    with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
+
+        def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
+            info = tarfile.TarInfo(name=name)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+        add_file_bytes("init_fn.onnx", init_fn.SerializeToString())
+        add_file_bytes("step_fn.onnx", step_fn.SerializeToString())
+        add_file_bytes("metadata.json", metadata.model_dump_json().encode("utf-8"))
+
+    buffer.seek(0)
+
+    return buffer.read()

kinfer/requirements.txt

@@ -1,9 +1,8 @@
 # requirements.txt
 
 # Machine Learning
-torch
 onnx
-onnxruntime
+onnxruntime==1.20.0
 
-# Protocol Buffers
-protobuf
+# Serialization
+pydantic

kinfer/rust/Cargo.toml

@@ -6,15 +6,17 @@ edition = "2021"
 
 [dependencies]
 
+async-trait = "0.1"
+flate2 = "1.0"
 futures-util = "0.3.30"
 ndarray = "0.16.1"
-ort = { version = "2.0.0-rc.6", features = [ "load-dynamic" ] }
-prost = "0.12"
+ort = { version = "2.0.0-rc.9", features = [ "load-dynamic" ] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+tar = "0.4"
 thiserror = "1.0"
+tokio = { version = "1.0", features = ["full"] }
 
 [dev-dependencies]
-rand = "0.8"
-
-[build-dependencies]
 
-prost-build = "0.12"
+rand = "0.8"

kinfer/rust/src/lib.rs

@@ -1,14 +1,5 @@
-pub mod kinfer_proto;
 pub mod model;
-pub mod onnx_serializer;
-pub mod serializer;
+pub mod runtime;
 
-pub use kinfer_proto::*;
 pub use model::*;
-pub use onnx_serializer::*;
-pub use serializer::*;
-
-#[cfg(test)]
-mod tests {
-    mod onnx_serializer_tests;
-}
+pub use runtime::*;