ONNX export of a k-Means

Download Python samples A Zip archive containing all samples can be found here: Samples of ONNX export

k-Means with Scikit-learn

from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# Generate data for clustering: num of samples, dimensions, num of clusters
X, y, centers = make_blobs(n_samples=50, n_features=3, centers=5, cluster_std=0.5, shuffle=True, random_state=42,  return_centers=True)
num_features = X.shape[1]
num_clusters = centers.shape[0]

# Define and train K-Means model
km = KMeans(n_clusters=num_clusters, init='k-means++', n_init=10, max_iter=500, tol=1e-04, random_state=42)
km.fit(X)

y_km = km.predict(X)

# Export model as ONNX
out_onnx = "kmeans.onnx"

initial_type = [('float_input', FloatTensorType([None, num_features]))]
onnx_model = convert_sklearn(km, initial_types=initial_type)

# for k-means models this meta info is mandatory. Otherwise convert process to TwinCAT-specific format will fail!
meta =  onnx_model.metadata_props.add()
meta.key = "sklearn_model"
meta.value = "KMeans"

with open(out_onnx, "wb") as f:
    f.write( onnx_model.SerializeToString())

According to our level of knowledge, only Scikit-learn with skl2onnx is currently capable of converting a k-Means to ONNX. For this reason, the description is limited to that.

ONNX Custom Attributes necessary The specification of the Custom Attribute "sklearn_model" and "KMeans" is necessary for k-means models, so that the conversion step in Beckhoff XML and Beckhoff BML works.