ONNX Inference

ONNX is an open format built to represent machine learning models. ONNX defines a common set of operators - the building blocks of machine learning and deep learning models - and a common file format to enable AI developers to use models with a variety of frameworks, tools, runtimes, and compilers.

Latest BrainFlow release adds an option to load some ONNX models into BrainFlow and use it from different programming languages. This tutorial is about how to use it and it’s current limitations.

Limitations

  • We’ve tested it only using scikit-learn and skl2onnx. You can try it with DL frameworks like Tensorflow, but maybe there will be some issues.
  • Input type can be only float or double
  • Disable ZipMap for classifiers
  • Max ONNX opset we tested so far is 11
  • We’ve added only default execution provider(CPU)

Example(training)

If you already have a model skip this part and go to inference example. Continue reading this section if you want to train one by yourself.

You can find full source code for training here.

Before running inference using BrainFlow API we need to train the model. In this example we will use scikit-learn. Dependencies we will need:

  • scikit-learn
  • brainflow
  • metric-learn
  • skl2onnx

You can install all of them from PYPI using python -m pip install %dependency name here%

Dataset collection is up to you, we will skip it in this tutorial. For features we will use band powers. To calculate band powers you can use get_avg_band_powers or new method get_custom_band_powers. Call them inside a moving window and create feature vectors and labels as below.

bands = DataFilter.get_avg_band_powers(data_in_window, eeg_channels, sampling_rate, True)
feature_vector = bands[0]
feature_vector = feature_vector.astype(float) # input type can be float or double, float is recommended because some ONNX operators work better with floats
dataset_x.append(feature_vector)
if data_type == first_class:
    dataset_y.append(0)
else:
    dataset_y.append(1)

It’s always good to take a look at the data and print some metrics, so you can run smth like this:

def print_dataset_info(data):
    x, y = data
    first_class_ids = [idx[0] for idx in enumerate(y) if idx[1] == 0]
    second_class_ids = [idx[0] for idx in enumerate(y) if idx[1] == 1]
    x_first_class = list()
    x_second_class = list()
    
    for i, x_data in enumerate(x):
        if i in first_class_ids:
            x_first_class.append(x_data.tolist())
        elif i in second_class_ids:
            x_second_class.append(x_data.tolist())
    second_class_np = np.array(x_second_class)
    first_class_np = np.array(x_first_class)

    logging.info('1st Class Dataset Info:')
    logging.info('Mean:')
    logging.info(np.mean(first_class_np, axis=0))
    logging.info('2nd Class Dataset Info:')
    logging.info('Mean:')
    logging.info(np.mean(second_class_np, axis=0))

This code prints mean values for band powers in both classes.

For this example we will use RandomForest as a classier.

def train_random_forest_mindfulness(data):
    model = RandomForestClassifier(class_weight='balanced', random_state=1, n_jobs=8, n_estimators=200)
    logging.info('#### Random Forest ####')
    scores = cross_val_score(model, data[0], data[1], cv=5, scoring='f1_macro', n_jobs=8)
    logging.info('f1 macro %s' % str(scores))
    model.fit(data[0], data[1])

    initial_type = [('mindfulness_input', FloatTensorType([1, 5]))]
    onx = convert_sklearn(model, initial_types=initial_type, target_opset=11, options={type(model): {'zipmap': False}})
    with open('forest_mindfulness.onnx', 'wb') as f:
        f.write(onx.SerializeToString())

Important: you have to disable ZipMap and specify target opset 11 or less

Snippet above generates onnx file, you should inspect the content of this file and check names of output nodes. Go to netron.app and load your model there.

This graph looks simple, for output node we will need probabilities instead labels.

Example(inference)

Now for demo purposes we will move to BrainFlow’s Java API. But you can do the same from any other supported programming language and even from game engines like Unity.

We will start with built-in default classifier for one of predefined metrics and after that I will show how ONNX classifier differs from the basic one.

Pair<double[], double[]> bands = DataFilter.get_avg_band_powers (data, eeg_channels, sampling_rate, true);
double[] feature_vector = bands.getLeft ();
BrainFlowModelParams model_params = new BrainFlowModelParams (BrainFlowMetrics.MINDFULNESS,
        BrainFlowClassifiers.DEFAULT_CLASSIFIER);
MLModel model = new MLModel (model_params);
model.prepare ();
double[] result = model.predict (feature_vector);
model.release ();

You can find this examples at BrainFlow docs.

Now, to switch default classifier to ONNX we need to change metric, classifier and provide extra options for model_params.

Pair<double[], double[]> bands = DataFilter.get_avg_band_powers (data, eeg_channels, sampling_rate, true);
double[] feature_vector = bands.getLeft (); // our model was trained using float type but BrainFlow can also handle double values
BrainFlowModelParams model_params = new BrainFlowModelParams (BrainFlowMetrics.USER_DEFINED,
        BrainFlowClassifiers.ONNX_CLASSIFIER);
model_params.file = "HERE YOU NEED TO WRITE THE FULL PATH TO YOUR ONNX MODEL";
model_params.output_name = "probabilities"; // its optional, brainflow can autodiscover some output nodes but its recommended to provide manually
MLModel model = new MLModel (model_params);
model.prepare ();
double[] result = model.predict (feature_vector);
model.release ();

There is a lot of room for improvements here, but it also unlocks some really great opportunities. Feel free to give it a try and report bugs if any!