Run this NB in Google Colab

ONNX model snapshots#

This notebook shows how to use the SaveToOnnx callback to periodically save the model of a running Session to ONNX format, and how to load the saved ONNX model back and use it to score new data.

[1]:
## Install and import the required libraries

This notebook requires onnxruntime. Install it if it isn’t already available:

pip install onnxruntime
[2]:
import numpy as np
import onnxruntime as rt
import os
from coniferest.datasets import single_outlier
from coniferest.isoforest import IsolationForest
from coniferest.session import Session
from coniferest.session.callback import Label, SaveToOnnx, TerminateAfter

Prepare a small dataset#

We use the single_outlier toy dataset that ships with coniferest.

[3]:
data, metadata = single_outlier()
data = data.astype(np.float32)

print(f"Data shape: {data.shape}")
Data shape: (10001, 2)

Run a session that saves the model to ONNX#

We attach SaveToOnnx as an on_decision_callback. Here we save every 5 decisions, keeping a numbered file for each save (overwrite=False), so we can inspect how the model evolves over time.

[4]:
ONNX_DIR = "onnx_models"
EXPERT_BUDGET = 20


def decision(index, x, session):
    # Non-interactive: pick a fixed answer just to drive the demo session
    return Label.REGULAR


model = IsolationForest(
    n_trees=100,
    random_seed=0,
)

session = Session(
    data=data,
    metadata=np.arange(len(data)),
    model=model,
    decision_callback=decision,
    on_decision_callbacks=[
        TerminateAfter(EXPERT_BUDGET),
        SaveToOnnx(
            directory=ONNX_DIR,
            filename="model.onnx",
            every_n_decisions=5,
            overwrite=False,
        ),
    ],
)

session.model.fit(data)
session.run()
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
[4]:
<coniferest.session.Session at 0x7856c9876660>

Let’s see which ONNX snapshots were saved during the session.

[5]:
saved_files = sorted(os.listdir(ONNX_DIR))
print(saved_files)
['model_10.onnx', 'model_15.onnx', 'model_20.onnx', 'model_5.onnx']

Load a saved ONNX model and use it#

Pick the last saved snapshot and load it with onnxruntime to run inference, without needing the original coniferest model object.

[6]:
last_snapshot_path = os.path.join(ONNX_DIR, saved_files[-1])

sess = rt.InferenceSession(last_snapshot_path)
input_name = sess.get_inputs()[0].name
label_name = "score"

onnx_scores = sess.run([label_name], {input_name: data})[0].reshape(-1)

print(f"Top 5 most anomalous scores (ONNX): {np.sort(onnx_scores)[:5]}")
2026-10-14 14:39:56.017836471 [W:onnxruntime:, graph.cc:125 MergeShapeInfo] Error merging shape info for output. 'score' source:{-1,1} target:{-1}. Falling back to lenient merge.
Top 5 most anomalous scores (ONNX): [-0.77463865 -0.7601073  -0.7487968  -0.74815553 -0.73900306]

Overwrite strategy example#

If you only care about the latest model and don’t want to keep every snapshot, use overwrite=True (the default) so the callback always writes to the same file.

[7]:
OVERWRITE_DIR = "onnx_model_latest"

callback = SaveToOnnx(directory=OVERWRITE_DIR, filename="model.onnx")

session_overwrite = Session(
    data=data,
    metadata=np.arange(len(data)),
    model=IsolationForest(n_trees=100, random_seed=0),
    decision_callback=decision,
    on_decision_callbacks=[
        TerminateAfter(EXPERT_BUDGET),
        callback,
    ],
)
session_overwrite.model.fit(data)
session_overwrite.run()

print(os.listdir(OVERWRITE_DIR))
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
The maximum opset needed by this model is only 8.
['model.onnx']

Saving the whole session with PickleSession#

Unlike SaveToOnnx, which only saves the model, PickleSession saves the entire session object (model, known labels, current candidate, scores, etc.) using Python’s pickle. This means an interrupted session can be fully resumed later, not just the model.

[8]:
from coniferest.session.callback import PickleSession
import pickle

PICKLE_DIR = "session_snapshots"

session_with_pickle = Session(
    data=data,
    metadata=np.arange(len(data)),
    model=IsolationForest(n_trees=100, random_seed=0),
    decision_callback=decision,
    on_decision_callbacks=[
        TerminateAfter(EXPERT_BUDGET),
        PickleSession(directory=PICKLE_DIR, every_n_decisions=5),
    ],
)
session_with_pickle.model.fit(data)
session_with_pickle.run()

print(os.listdir(PICKLE_DIR))
['session.pickle']

Now let’s load the saved session back and confirm it’s fully resumed, including the known labels.

[9]:
saved_sessions = sorted(os.listdir(PICKLE_DIR))
last_session_path = os.path.join(PICKLE_DIR, saved_sessions[-1])

with open(last_session_path, "rb") as f:
    restored_session = pickle.load(f)

print(f"Known labels: {restored_session.known_labels}")
print(f"Model type: {type(restored_session.model)}")
Known labels: {np.uint64(5054): <Label.REGULAR: 1>, np.uint64(10000): <Label.REGULAR: 1>, np.uint64(5619): <Label.REGULAR: 1>, np.uint64(547): <Label.REGULAR: 1>, np.uint64(2956): <Label.REGULAR: 1>, np.uint64(2710): <Label.REGULAR: 1>, np.uint64(4736): <Label.REGULAR: 1>, np.uint64(4164): <Label.REGULAR: 1>, np.uint64(2606): <Label.REGULAR: 1>, np.uint64(6836): <Label.REGULAR: 1>, np.uint64(6416): <Label.REGULAR: 1>, np.uint64(3016): <Label.REGULAR: 1>, np.uint64(3124): <Label.REGULAR: 1>, np.uint64(7216): <Label.REGULAR: 1>, np.uint64(2508): <Label.REGULAR: 1>, np.uint64(3688): <Label.REGULAR: 1>, np.uint64(9860): <Label.REGULAR: 1>, np.uint64(7008): <Label.REGULAR: 1>, np.uint64(5238): <Label.REGULAR: 1>, np.uint64(5152): <Label.REGULAR: 1>}
Model type: <class 'coniferest.isoforest.IsolationForest'>