Model Export & Serving Runtime API Reference

k3_node.export.onnx_exporter.export_onnx(model_or_task, output_path, dummy_inputs=None, opset=17, dynamic_axes=True, input_names=None, output_names=None, verbose=False)

Exports a K3-Node GNN model or task to high-performance ONNX format.

Supports arbitrary Graph Neural Networks (GCN, GAT, GraphSAGE, GIN, SchNet, materials models, and task estimators) with dynamic graph sizing (varying numbers of nodes and edges).

Parameters:

Name	Type	Description	Default
model_or_task	Any	A K3-Node task instance (e.g. NodeClassifier, GraphClassifier) or model instance (e.g. GCN, SchNet, CHGNet).	required
output_path	Union[str, Path]	Target path for the .onnx file.	required
dummy_inputs	Optional[Any]	Optional sample input data (e.g., PyG Data object, tuple of tensors). If None, automatically generated based on model topology.	None
opset	int	ONNX operator set version. (default: 17)	17
dynamic_axes	bool	Whether node and edge dimensions should be dynamic. (default: True)	True
input_names	Optional[List[str]]	Optional custom names for input tensors.	None
output_names	Optional[List[str]]	Optional custom names for output tensors.	None
verbose	bool	Whether to print verbose export progress. (default: False)	False

Returns:

Type	Description
Path	Path object pointing to the generated .onnx file.

k3_node.export.tflite_exporter.export_tflite(model_or_task, output_path, dummy_inputs=None, quantization=None, representative_dataset=None, verbose=False)

Exports a K3-Node GNN model or task to an optimized TensorFlow Lite flatbuffer.

Supports float32, float16 (FP16), and dynamic range INT8 quantization for low-latency deployment on edge, mobile, and embedded hardware.

Parameters:

Name	Type	Description	Default
model_or_task	Any	A K3-Node task instance or model instance.	required
output_path	Union[str, Path]	Target path for the .tflite file.	required
dummy_inputs	Optional[Any]	Optional sample input data (e.g., PyG Data object, tuple of tensors).	None
quantization	Optional[str]	Quantization strategy: - None or "fp32": Full 32-bit floating point precision. - "fp16": 16-bit floating point quantization (~2x smaller & faster on GPU/NPU). - "int8_dynamic": Dynamic range integer quantization (~4x smaller with 8-bit weights). - "int8_full": Full 8-bit integer quantization (requires representative_dataset).	None
representative_dataset	Optional[Callable[[], Generator[List[ndarray], None, None]]]	Generator of representative calibration inputs for "int8_full".	None
verbose	bool	Whether to print verbose progress. (default: False)	False

Returns:

Type	Description
Path	Path object pointing to the generated .tflite file.

k3_node.export.tensorrt_exporter.export_tensorrt(model_or_task_or_onnx, output_path, dummy_inputs=None, precision='fp16', workspace_gb=1, min_shapes=None, opt_shapes=None, max_shapes=None, verbose=False)

Compiles a K3-Node model or ONNX file into an ultra-low latency NVIDIA TensorRT engine.

Uses the Python tensorrt API if available on the GPU host, or invokes NVIDIA's trtexec binary directly.

Parameters:

Name	Type	Description	Default
model_or_task_or_onnx	Any	K3-Node model/task, or an existing .onnx file path.	required
output_path	Union[str, Path]	Target path for the compiled .engine or .plan binary.	required
dummy_inputs	Optional[Any]	Optional input sample to determine shapes and topology.	None
precision	str	Precision mode: "fp32", "fp16", or "int8". (default: "fp16")	'fp16'
workspace_gb	int	Max GPU memory in GB allocated for TensorRT engine building. (default: 1)	1
min_shapes	Optional[Dict[str, Tuple[int, ...]]]	Optional minimum dynamic shapes dictionary (e.g. {'x': (1, 16)}).	None
opt_shapes	Optional[Dict[str, Tuple[int, ...]]]	Optional optimal dynamic shapes dictionary (e.g. {'x': (100, 16)}).	None
max_shapes	Optional[Dict[str, Tuple[int, ...]]]	Optional maximum dynamic shapes dictionary (e.g. {'x': (10000, 16)}).	None
verbose	bool	Whether to log engine building progress. (default: False)	False

Returns:

Type	Description
Path	Path object pointing to the compiled TensorRT .engine file.

k3_node.export.tensorrt_exporter.generate_triton_config(model_name, output_dir, in_channels, out_channels, backend='onnxruntime', max_batch_size=0)

Generates a complete Triton Inference Server model repository directory structure and config.pbtxt.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the model in Triton (e.g., 'cora_gcn').	required
output_dir	Union[str, Path]	Root path for the model repository folder.	required
in_channels	int	Number of node input features.	required
out_channels	int	Number of output features or classes.	required
backend	str	Backend engine ('onnxruntime' or 'tensorrt_plan').	'onnxruntime'
max_batch_size	int	Maximum batch size (default: 0 for dynamic graph sizing).	0

Returns:

Type	Description
Path	Path to the generated config.pbtxt.

k3_node.export.runtime.ONNXModel

High-performance serving wrapper for exported ONNX GNN models.

Requires only onnxruntime and numpy. Completely decoupled from Keras, PyTorch, and TensorFlow for lightweight production microservices.

Example

from k3_node.export import ONNXModel
model = ONNXModel("cora_gcn.onnx")
preds = model.predict(graph_data)

__init__(model_path, providers=None, session_options=None)

Initializes the ONNX runtime inference session.

Parameters:

Name	Type	Description	Default
model_path	Union[str, Path]	Path to the .onnx model file.	required
providers	Optional[List[str]]	Execution providers list (e.g. ['CUDAExecutionProvider', 'CPUExecutionProvider']). If None, automatically picks the fastest available provider.	None
session_options	Optional[Any]	Optional custom ONNX Runtime SessionOptions.	None

predict(data=None, *args, **kwargs)

Runs low-latency inference on graph data.

Accepts

• PyG / K3-Node Data object
• Dictionary of tensors
• Positional numpy arrays

Parameters:

Name	Type	Description	Default
data	Any	Input graph Data, dictionary, or array.	None
*args	Any	Additional positional inputs.	()

Returns:

Type	Description
ndarray	Numpy array containing model predictions or logits.

k3_node.export.runtime.TFLiteModel

Lightweight serving wrapper for TensorFlow Lite flatbuffer GNN models.

Requires only standard tensorflow or tflite_runtime. Ideal for mobile, Raspberry Pi, and edge embedded devices.

Example

from k3_node.export import TFLiteModel
model = TFLiteModel("model.tflite")
preds = model.predict(graph_data)

__init__(model_path)

Initializes the TFLite interpreter.

Parameters:

Name	Type	Description	Default
model_path	Union[str, Path]	Path to the .tflite model file.	required

predict(data=None, *args, **kwargs)

Runs inference using the TFLite interpreter.

Parameters:

Name	Type	Description	Default
data	Any	Input graph Data, dict, or numpy array.	None
*args	Any	Additional positional inputs.	()

Returns:

Type	Description
ndarray	Numpy array containing prediction results.