Add an Existing Script to the MLTK

The MLTK is a collection of tools to aid the development of machine learning models for Silicon Labs embedded devices. As part of its offering, the MLTK features an optional Python API to simplify ML model development using the Tensorflow framework. The MLTK Python API combines and simplifies several different ML model development tasks including:

The MLTK Python API is effectively a light wrapper around the Tensorflow API. As such, very little needs to be done to convert an existing Tensorflow model training script to support the MLTK, and, thus, gain the additional MLTK features (quantization, evaluation, profiling).

NOTE: Converting your existing model script to the MLTK is not required. The only hard requirement to use your model on a Silicon Labs embedded device is that your model be converted to the .tflite format with int8 quantization.

Setup

Before continuing with this tutorial, ensure you have the MLTK Python package installed and that the virtual environment is activated (if applicable).

Existing Script

For this tutorial, we use the Keras example: Simple MNIST convnet

The source code for the model training script may be found on GitHub.

It has also been copy & pasted as follows:

Original mnist_convnet.py

import numpy as np
from tensorflow import keras
from tensorflow.keras import layers

"""
## Prepare the data
"""

# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

# Load the data and split it between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")


# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

"""
## Build the model
"""

model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
)

model.summary()

"""
## Train the model
"""

batch_size = 128
epochs = 15

model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)

"""
## Evaluate the trained model
"""

score = model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])

MLTK Modifications

Download the mnist_convnet.py script to your local PC, and add the following modifications to it so that it we can execute it in the MLTK.

All modifications are prefixed with the comment: # Modified for the MLTK

Copy & paste the following into the mnist_convnet.py script on your local PC:

Modified mnist_convnet.py

import numpy as np
from tensorflow import keras
from tensorflow.keras import layers

"""
## Prepare the data
"""

# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

# Load the data and split it between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")


# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

"""
## Build the model
"""

model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
)

model.summary()

"""
## Train the model
"""

batch_size = 128
epochs = 15

model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

# Modified for the MLTK
# We only want to train when calling the train_model() API
#model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)

"""
## Evaluate the trained model
"""

# Modified for the MLTK
# We only want to evaluate when calling teh evaluate_model() API
#score = model.evaluate(x_test, y_test, verbose=0)
#print("Test loss:", score[0])
#print("Test accuracy:", score[1])


# Modified for the MLTK
# The follow allows for running this existing model training script in the MLTK

# Import the MLTK Python package
import mltk.core as mltk_core 
import tensorflow as tf

# Instantiate the MltkModel object
class MyModel(
    mltk_core.MltkModel,
    mltk_core.TrainMixin,
    mltk_core.DatasetMixin,
    mltk_core.EvaluateClassifierMixin, 
):
    pass
my_model = MyModel()

# This simply returns the previously built KerasModel
def my_model_builder(mltk_model):
    return model

# This callback allows for loading the training or testing dataset
def my_dataset_loader(
    subset:str,
    test:bool, 
    **kwargs
):
    if subset == 'training':
        return x_train, y_train 
    else:
        return x_test, y_test

# This is used by the TfliteConverter when generating the quantized .tflite
def my_representative_dataset_generator():
    for input_value in tf.data.Dataset.from_tensor_slices(my_model.x).batch(1).take(100):
        yield [input_value]


# Populate the various MltkModel properties
my_model.build_model_function = my_model_builder
my_model.dataset = my_dataset_loader
my_model.epochs = epochs
my_model.batch_size = batch_size 
my_model.validation_split = 0.1
my_model.classes = ['0','1','2','3','4','5','6','7','8','9']
my_model.tflite_converter['representative_dataset'] = my_representative_dataset_generator

# If this script is being invoked from the command-line then execute it now, e.g.:
# python mnist_convnet.py
if __name__ == '__main__':
    from mltk import cli
    cli.get_logger(verbose=False)

    train_results = mltk_core.train_model(my_model, clean=True)
    print(train_results)

    tflite_eval_results = mltk_core.evaluate_model(my_model, verbose=True)
    print(tflite_eval_results)

    profiling_results = mltk_core.profile_model(my_model)
    print(profiling_results)

Run the modified script

If you invoke the modified script with the command:

python mnist_convnet.py

Then, instead of just using Tensorflow for model training like the original script used, the modified script will use the MLTK and perform the following:

  1. Train the model using Tensorflow

  2. Generate a quantized .tflite model file

  3. Evaluate the quantized .tflite model

  4. Profile the .tflite model in the Tensorflow-Lite Micro simulator

Run from command-line

Additionally, the model script can either directly invoke the model script from Python (e.g. python mnist_convnet.py ) or use the mltk command, e.g.:

# Train the model using the "mltk train" command
mltk train mnist_convnet.py
# Evaluate the model using the "mltk evaluate" command
mltk evaluate mnist_convnet.py
# Profile the model using the "mltk profile" command
mltk profile mnist_convnet.py

Next Steps

This tutorial demonstrates how minimal modifications can be added to an existing model training script so that it can run in the MLTK.

For complete, non-trivial models, please see the Reference Models that come with the MLTK.