Machine Learning at the Edge with Scikit-Learn, Keras, BentoML, and HiveMQ

Machine learning at the edge is becoming a necessity for many IoT and IIoT use cases that require quick decisions without relying on cloud connectivity. However, it's not always easy to pull off. With limited processing power on edge devices, spotty network connections, and the constant headache of keeping ML models up-to-date, it’s enough to make any developer want to pull their hair out.

But don't worry, in this article, we're going to show you how to tackle these challenges head-on using these powerful tools: Scikit-Learn, BentoML and HiveMQ's MQTT broker.

BentoML is like a Swiss Army knife for ML deployment. It helps you package up your Tensorflow models, making them easy to ship out to edge devices. Add HiveMQ, the most reliable MQTT platform, and you have a solution for making sure your data gets where it needs to go.

We'll walk you through the details of setting this up to get your ML models running smoothly at the edge, even when the Internet is acting up. So, grab a coffee, fire up your favorite code editor, and let's dive in. It's time to bring some machine learning muscle to the edge of your network.

Building a Predictive Quality Model for Manufacturing

In this project, we'll develop a machine learning model based on TensorFlow designed to predict a quality index for manufacturing batches. Our model will utilize telemetry data collected throughout the production process. The aim is to generate a quality prediction at the conclusion of each batch. This predictive capability will enable manufacturers to anticipate potential quality issues, optimize their processes, and make data-driven decisions to improve overall production efficiency. This quality prediction will then be published in the Unified Namespace to be consumed by other systems.

The end to end architecture will look like the diagram below.

Prerequisites

In order to build and run this example, you will need the following :

• Python 3.9.6 minimum

• BentoML (Follow the documentation to install it)

• Visual Studio Code (or your favorite Jupiter Notebook editor)

• Intellij IDEA for the HiveMQ extension creation

• Java OpenJDK 21 or earlier

• HiveMQ Edge version 2024.7 minimum (Also works with HiveMQ Enterprise 4.34 minimum)

Everything can be executed on the same machine if you want to build a lab. I will not give details on how to configure the bridge on the edge to send the telemetry and inference result to the central broker holding the UNS as this is pretty standard and available in our documentation.

You can find the content of the article in this GitHub repository.

Build and Train the ML Model

Create a Python Virtual Environment

It’s always a good idea to isolate the code and dependencies in a virtual environment to avoid side effects with your existing system.

Once created, activate it.

Install Python Dependencies

In order to build and run our model, we will need some libraries.

• BentoML: Our Swiss Army knife for ML

• Scikit-Learn: An end-to-end platform for machine learning

• Keras: Consistent API for modeling your TensorFlow models

Install all at once (if you don’t want to use the requirements.txt shared in the GitHub repository).

Prepare Your Data

Usually, when you train your model, you do it on your own data that you’ve collected during days, weeks, or even more. In our demo, since I don’t have a factory generating data, I will use this dataset to train my model. It includes 17 values from temperature sensors from different roasting machines. The values are captured every minute over a long period.

After you download the dataset, create a folder called “ProductionQuality” and drop the files there.

Create a file called “train.py” and add the following sections (GitHub).

The script involves data preprocessing, training a neural network model, and then saving and using the model for predictions.

1. Imports:

   • bentoml: Library used for packaging and serving machine learning models.

   • pandas: Library for data manipulation and analysis.

   • os: Module for interacting with the operating system.

   • StandardScaler and train_test_split from sklearn: Used for data normalization and splitting data into training and test sets.

   • Sequential, Dense, Dropout, RMSprop from keras/tensorflow: Used for building and training a neural network model.

2. File Checking and Data importing:

   • Lists and prints the files in the './ProductionQuality/' directory to ensure necessary data files are present.

   • Reads CSV files containing the input data, the output quality indexes for training, and a sample submission file for testing.

   • Prints the shape of the imported datasets.

3. Data Merging:

   • Merges x_data with y_train to create train_df and with Y_submit to create test_df, using date_time as the key.

   • Removes the 'quality' column from test_df, as it is not part of the test features.

4. Data Integrity Check and Data Cleaning:

   • Asserts that the row count of train_df matches y_train and that of test_df matches Y_submit to ensure proper merging.

   • Separates the target variable y (quality) from the features in train_df.

   • Drops the 'date_time' column from the datasets since it's not used as a feature.

5. Dataset Splitting and Scaling:

   • Splits train_df into X_train, X_test, and splits y into y_train, y_test using a 33% test size.

   • Scales the features using StandardScaler to normalize the data.

6. Model Building and Compilation:

   • Creates a Sequential neural network model with several layers:

     • Input layer: Dense with tanh activation.

     • Hidden layers: Dense with relu activation and Dropout for regularization.

     • Output layer: Dense with a linear activation function for regression tasks.

   • Compiles the model specifying mean squared error as the loss function, RMSprop as the optimizer, and metrics for evaluation.

7. Model Training and Model Evaluation:

   • Fits the model on X_train and y_train with specified batch_size and epochs.

   • Includes validation on X_test and y_test.

   • Evaluates the model on the test set and prints the test mean squared error (mse) and mean absolute error (mae).

8. Model Saving and Predictions tests:

   • Saves the trained model using BentoML for deployment, assigning it the tag 'qualitycheck'.

   • Uses the trained model to predict quality for the test_df and saves these predictions into 'submission.csv'.

Run the program to create your machine learning model and save it with BentoML.

You should get a Bento model tag at the end of the execution, keep it for the next step.

Now that we have a model for predicting the quality index based on telemetry information from the production batch, let's move on to preparing it for production operation.

Pack and Deploy the Model

In order to run the model in production, we need to follow a few steps required by BentoML to pack and serve the model for inference.

• Create a service

• Describe your Bento properties

• Test our Bento

• Build the Docker image with your Bento

Create your Service

Create a folder called “service” with a new file called service.py that has the following content (GitHub).

This code (Python) defines a BentoML service that wraps a machine learning model for making predictions.

BENTO_MODEL_TAG is an important part here, you need to update it with the Tag generated by the train script.

Describe your Bento

Create a new file called “bentofile.yaml” and insert the following code (GitHub).

The file is used by BentoML to prepare the Bento that will include the necessary resources to run our model and serve an inference API. We include all our python files (in our case the above service.py) and indicate some dependencies.

Test your Bento

Once you’ve defined the service and the properties of your Bento, you can test it with the below command:

At this point you should be able to access the Swagger on http://localhost:3000 and test the inference.

Here is a sample input data that you can use to test it directly in Swagger:

If everything is working fine, you should get the quality index in the Body of the response. Note that the more you use the inference API the faster it will answer, here in the test the first execution took 48ms and after it we drop to 13ms which is pretty good considering that we run on a CPU and not on a GPU.

Build Your Bento

Run the following command in the “service” directory to build your Bento.

You should get a similar output.

Pack it as a Docker Image

The last step is to prepare the Bento to be served from a docker container, as we already have it running we only need to run the following command to create the image (change the Bento ID with yours).

You can then test your newly created image by running:

In a similar way to out-of-container execution, you can access the Swagger and perform an inference test.

Deploy the ML Community Extension

In order to save some time in building an extension that captures messages on MQTT Topic to send them to your BentoML inference endpoint, I’ve created a Community extension to do the job. You can download it on GitHub.

Once you have it, unzip it in the “extensions” folder. Enter the extension folder and rename the “config-template.json” to “config.json”.

Edit the content to reflect your inference flow. Mine looks like the following:

Once you are ready, start your broker and check that you get the following message in the logs:

If not, check the error message to find what’s wrong.

Connect the Dots

Now that we have all the components of our architecture in place, we can run a test.

Use your favorite MQTT client to send the following message (or the one adapted to your model) in the MQTT topic “test/inputs.”

You should see on the destination topic (in my test “test/prediction”) the result of the inference being published.

Well done, you've performed your first real-time inference based on MQTT.

Wrap Up

We’ve covered a lot of ground on running ML at the edge with Scikit-Learn, BentoML, and HiveMQ. Here's a quick summary:

• Scikit-Learn: One of the go-tos for building lean, mean ML models that will run smoothly on your edge devices.

• BentoML: The secret sauce for packaging those models up and shipping them out without losing your mind.

• HiveMQ MQTT: Your reliable MQTT platform, keeping everything talking smoothly even when the network gets cranky.

This example isn't just for show, it’s about solving real problems out in the wild. If you're trying to make factories smarter, we’ve only scratched the surface. Edge technology is moving fast, and what's hot today might be old news tomorrow. You can explore alternatives for the ML framework like Pytorch or Scikit-Learn and other packaging solutions. So keep your eyes open and your skills sharp.

If you take anything away from this, it's this: the devil is in the details. Optimize those models, handle your data like it's gold, and plan for when things inevitably go sideways. Now it's your turn. Take these tools, go forth, and build something awesome at the edge. Who knows? You might just cook up the next big thing.