Note: This project derived from this blog post here/
This code can be run on the Teensy 4.0 board and test the MNIST TensorFlow Lite model
I'm using for evaluation and testing purposes in my Machine Learning for Embedded
post series which starts from here.
The most relevant post is this one
as it's the same code running on the STM32F746.
First you need to build and upload the project on the stm32f7. To do that follow the instructions in the build section. After that you can use the jupyter notebook to hand-draw a digit and then upload the digit on the stm32f7 and get the prediction back. Please follow the guide inside the notebook.
In order to run the notebook, you need python3, tensorflow and PySerial. I've used Ubuntu 18.04 and miniconda, but conda is not really needed. In any case it's good to run the following commads on a virtual environment.
Example for conda
conda create -n nn-env python
conda activate nn-env
conda install -c conda-forge numpy
conda install -c conda-forge jupyter
conda install -c conda-forge tensorflow-gpu
jupyter notebook
And then browse to the jupyter_notebook/MNIST-TensorFlow.ipynb
and run/use the notebook.
Note: Personally for Jupyter notebooks I'm using the
nvcr.io/nvidia/tensorflow:20.03-tf2-py3
docker image which can be found here.
The code implements a debug port that you can use to run commands from the terminal. This is the pinout for Teensy 4.0
UART | Tx | Rx |
---|---|---|
UART2 | 14 | 15 |
To fetch the repo use either of those two commands:
git clone [email protected]:dimtass/imxrt1062-tflite-micro-mnist.git
# or for http
git clone https://[email protected]/dimtass/imxrt1062-tflite-micro-mnist.git
There are a few options that you can use to build the code. To build the code you need
CMAKE and GCC toolchain installed to you system, but you can also use docker to build
the code if you want to skip installing stuff to your host. If you already have CMAKE
then in the build.sh
file you need to set the path of your toolchain in TOOLCHAIN_DIR
.
Then to build the code using the default settings run this command:
./build.sh
The available options are:
CLEANBUILD
: set totrue
to delete the build folder and re-built. Default value isfalse
which means faster builds when you do changes in the code.USE_CMSIS_NN
: This flag selects which kernels are used in the TF-Lite micro. Default value isOFF
which means that the default kernels are used. If it's enabled (set toON
) then the cmsis-nn kernels insource/libs/tensorflow/tensorflow/lite/micro/kernels/cmsis-nn
are used.USE_COMP_MODEL
: This option selects which model is used. Default value isOFF
which means that the uncompressed model is used (float16_t model and weights). If you set it toON
then it uses a model with compressed weights.
If you want to have the same build environment like the one I've used, then you can use my CDE image for STM32 (don't mind that) and docker like this:
./docker-build.sh "USE_COMP_MODEL=ON USE_CMSIS_NN=OFF ./build.sh"
CLEANBUILD=true USE_COMP_MODEL=ON USE_CMSIS_NN=OFF ./build.sh
CLEANBUILD=true USE_COMP_MODEL=OFF USE_CMSIS_NN=ON ./build.sh
The easiest way to flash the firmware is to use the teensy_loader_cli
tool from the console.
To do that you need to cd
to the repo directory then run one of the build commands and if
it's successful then use this command to flash the firmware:
teensy_loader_cli -v -w --mcu=imxrt1062 build-imxrt/flexspi_nor_release/imxrt1062-tflite-micro-mnist.hex
Then you should see this message:
Teensy Loader, Command Line, Version 2.1
Read "build-imxrt/flexspi_nor_release/imxrt1062-tflite-micro-mnist.hex": 367472 bytes, 18.1% usage
Waiting for Teensy device...
(hint: press the reset button)
Then you need to press the reset button on the Teensy 4.0 board and the flashing will start:
Found HalfKay Bootloader
Read "build-imxrt/flexspi_nor_release/imxrt1062-tflite-micro-mnist.hex": 367472 bytes, 18.1% usage
Programming.........................................................................................................................................................................................................................................................................................................................................................................
Booting
The above example is when the FW is built with this command:
CLEANBUILD=true USE_COMP_MODEL=ON USE_CMSIS_NN=OFF ./build.sh
There are two commands supported using the UART port. You can connect a USB-to-UART to Teensy using this pinout:
USB-to-UART | Teensy 4.0 |
---|---|
Tx | 14 |
Rx | 15 |
GND | G |
Then you can use a terminal (I'm using CuteCom
) and just send one of these commands, folowed
by a \n
or \r\n
if you like, but it's not really needed:
1
: Executes theViewModel()
function that displays information about the model. This is a sample output:
Model input:
dims->size: 4
dims->data[0]: 1
dims->data[1]: 28
dims->data[2]: 28
dims->data[3]: 1
input->type: 1
Model output:
dims->size: 2
dims->data[0]: 1
dims->data[1]: 10
2
: Runs the model inference. This is a sample output:
Running inference...
DEPTHWISE_CONV_2D: 6.283290 msec
MAX_POOL_2D: 0.829997 msec
CONV_2D: 165.644699 msec
MAX_POOL_2D: 0.246270 msec
CONV_2D: 25.556931 msec
FULLY_CONNECTED: 0.781590 msec
FULLY_CONNECTED: 0.098000 msec
SOFTMAX: 0.021400 msec
Done in 199.462173 msec...
Out[0]: 0
Out[1]: 0
Out[2]: 0
Out[3]: 0
Out[4]: 0
Out[5]: 0
Out[6]: 0
Out[7]: 0
Out[8]: 1.000000
Out[9]: 0
CMSIS version
: 5.0.4CMSIS-NN version
: V.2.0.0CMSIS-DSP version
: V1.7.0NXP SDK version
: 1.2
My code's license is MIT and you can use the code however you like. Used libraries may have other licenses!
Dimitris Tassopoulos [email protected]