How to Add an ONNX Model

Use ONNX for finite inference or repeated inference inside the same communication-step world as mechanistic modules.

Step 1: export the model

PyTorch:

import torch
 
model = MyClassifier()
model.eval()
dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["probabilities"],
)

Step 2: wrap it as a BioModule

import numpy as np
import biosimulant as biosim
 
 
class ONNXClassifier(biosim.BioModule):
    execution_policy = biosim.ExecutionPolicy.ONCE_BEFORE_RUN
 
    def __init__(self, model_path: str = "data/assets/model.onnx"):
        import onnxruntime as ort
 
        self.session = ort.InferenceSession(model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
 
    def inputs(self):
        return {
            "state_vector": biosim.SignalSpec.array(
                dtype="float32",
                shape=(10,),
            )
        }
 
    def outputs(self):
        return {
            "classification": biosim.SignalSpec.scalar(dtype="int64"),
            "probabilities": biosim.SignalSpec.array(dtype="float32", shape=(3,)),
        }
 
    def execute(self, inputs, *, context):
        input_data = np.asarray(inputs["state_vector"].value, dtype=np.float32).reshape(1, -1)
        result = self.session.run([self.output_name], {self.input_name: input_data})
        prediction = np.asarray(result[0][0], dtype=np.float32)
        class_idx = int(np.argmax(prediction))
        return {
            "classification": class_idx,
            "probabilities": prediction.tolist(),
        }

Step 3: declare the manifest

schema_version: "2.0"
title: "Neuron State Classifier"
description: "ONNX classifier for neuron state vectors"
standard: onnx
 
biosim:
  entrypoint: "src.onnx_classifier:ONNXClassifier"
  init_kwargs:
    model_path: "data/assets/model.onnx"
  communication_step: 0.001
 
io:
  inputs:
    - name: state_vector
      signal_type: array
      dtype: float32
      shape: [10]
  outputs:
    - name: classification
      signal_type: scalar
      dtype: int64
    - name: probabilities
      signal_type: array
      dtype: float32
      shape: [3]
 
onnx:
  task: classification
  model_file: data/assets/model.onnx
  inputs:
    - name: input
      dtype: float32
      shape: [-1, 10]
  outputs:
    - name: probabilities
      dtype: float32
      shape: [-1, 3]
 
runtime:
  python_version: "3.12"
  dependencies:
    packages:
      - onnxruntime>=1.16
      - numpy>=1.24

Step 4: compose it in a lab

models:
  - package: biosimulant/neuro-mechanistic
    version: 1.0.0
    alias: neuron
 
  - path: ../models/state-classifier
    alias: classifier
 
wiring:
  - from: neuron.state_vector
    to:
      - classifier.state_vector
 
runtime:
  duration: 0.05
  communication_step: 0.001

Keep ONCE_BEFORE_RUN for finite classification or a pure inference pipeline. Use ONCE_AFTER_RUN for final analysis, or EACH_WINDOW when the ONNX model must consume evolving simulation state. Lab and model manifests do not add an execution-mode field.

Next steps

Wrap an External Simulator Package & Publish