Training a Custom Model for a Kendryte K210 Board - Key Insights

For some time, I’ve been deeply engrossed in the challenge of training a custom model for my Maix Dock M1, a K210 board. This journey hasn’t been without its frustrations, especially due to the lack of comprehensive documentation, which turned out to be a significant hurdle. However, after a lot of trial and error, I’m excited to share that I’ve finally figured out a reliable approach to get things working.

In this article, I’ll walk you through some of the key insights and tips I’ve gathered along the way, which could save you a lot of time and effort if you’re on a similar path.

The Struggle with Documentation

The K210 boards equipped with its Kendryte Processing Unit (KPU), are a fascinating platform for edge AI applications. However, my initial excitement was quickly tempered by the realization that the documentation available was not as thorough as I had hoped. This lack of clear guidance meant that I spent countless hours experimenting, researching, and occasionally stumbling upon solutions by chance or through community forums.

Cracking the Code

After much perseverance, I managed to overcome the initial challenges. The breakthrough was particularly rewarding when I successfully ran a custom Fashion MNIST classifier on the K210 board. Seeing the model classify images on the device was a testament to the hard work and patience invested.

Key Tips for Training Models on K210

Here are some crucial tips that I learned during this process:

1. Use nncase v0.2.0 Beta4

One of the most significant lessons I learned was the importance of using the right tools. Initially, I spent a considerable amount of time trying to get things working with the latest version of nncase, a neural network compiler for Kendryte K210. Unfortunately, I kept running into issues that I couldn’t resolve on my own. It wasn’t until I received help on GitHub and I realized the latest version wasn’t compatible with my needs.

The solution was to use nncase v0.2 beta. If you’re trying to get a model working on the K210, I highly recommend sticking with this version. You can download nncase v0.2 beta from this link.

This is how you convert from tflite to kmodel

./ncc compile /workspace/myfolder/Maix_Toolbox/fashionmnist2.tflite /workspace/myfolder/Maix_Toolbox/fashionmnist2.kmodel -i tflite -o kmodel -t k210 --dataset /workspace/myfolder/Maix_Toolbox/images 

A successful run should show output like this:

1. Import graph...
2. Optimize Pass 1...
3. Optimize Pass 2...
4. Quantize...
4.1. Add quantization checkpoints...
4.2. Get activation ranges...
Plan buffers...
Run calibration...
[==================================================] 100% 0.163s
4.5. Quantize graph...
5. Lowering...
6. Optimize Pass 3...
7. Generate code...
Plan buffers...
Emit code...
Working memory usage: 62720 B

SUMMARY
INPUTS
0 Input_0 1x1x32x32
OUTPUTS
0 StatefulPartitionedCall_1:0 1x10

You can even run inference using newly created kmodel by using this command:

./ncc infer /workspace/myfolder/Maix_Toolbox/fashionmnist2.kmodel  .  --dataset  /workspace/myfolder/Maix_Toolbox/images/ 

In both the commands the dataset parameter points to a directory that has images used for post quantization. Make sure to use a hex viewer for opening *.bin files created by the infer command.

2. Adopt Code to Convert TensorFlow Model to TFLite

As mentioned earlier, the conversion from an .h5 file to a .tflite file is broken in the current version of TensorFlow (2.16). To bypass this issue, you’ll need to adopt a specific code snippet that ensures the conversion process works as expected. Here’s how you can do it:

pip install -U tf_keras
import tensorflow as tf

model = tf.keras.models.load_model('your_model.h5')
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

with open('model.tflite', 'wb') as f:
f.write(tflite_model)

This hack allows you to get around the current issues with the TensorFlow conversion process, ensuring that your model is ready for deployment on the K210.

3. Align Padding, Stride, and Pooling

To ensure that your model runs smoothly on the K210, it’s crucial to adopt padding, stride, and pooling configurations that align with the examples provided in this repository. The K210’s KPU has specific requirements, and adhering to these configurations can prevent a lot of headaches. Here’s a quick reference:

  • Padding: Use same for AveragePooling2D and valid for Conv2D, depending on the layer.
  • Stride: strides = (2,2),
  • Pooling: pool_size = (2,2)

By sticking to these settings, your model should work as expected on the K210.

4. Current Limitations and Next Steps

So far, I’ve only been able to get things working for grayscale images. This is a significant milestone, but the next step is to extend this functionality to RGB images. I’m optimistic that the process won’t be drastically different, and I’m eager to see how the K210 handles more complex image data.

Conclusion

Training a custom model for the K210 board is a rewarding yet challenging endeavor. The lack of clear documentation and some software incompatibilities were the major obstacles in my journey. However, by sticking with nncase v0.2 beta, adopting the TensorFlow conversion hack, and aligning your model’s architecture with official guidelines, you can successfully deploy models on the K210.

If you’re looking to embark on a similar project, I hope these tips help you avoid some of the pitfalls I encountered. With persistence and the right tools, the K210 board can be a powerful platform for your AI projects. Stay tuned as I explore the next phase of working with RGB images!