Number Predictor using Tensorflow.js and MNIST dataset

I recently starting learning Tensorflow.js to create machine learning models on the browser. It took me a few weeks to learn the API but I finally got something working.

The demo is a simple number predictor demo built using a Tensorflow.js model trained on the MNIST dataset.

There are two models in this example. One that was already trained locally on my computer and the browser model which is train right on your browser. I was very impressed by the accuracy of the browser trained model. I was able to get about 95% accuracy in some cases. I might create a tutorial but for now here's a link to the demo and source code.

Model Architecture

A Convolutional Neural Network (CNN) with 5 layers, 1 input 3 hidden and 1 for output. See code for more on activation function and other info on the model.

General Architecture

1. Load the MNIST dataset if training the model in the browser.
2. Load the pre-trained model or the browser trained model.
3. Take the input from the canvas element and do some processing to the pixels.
4. Take the input and run it throw the trained model.
5. Take the predictions and determine which one has the higher probability.
6. Display the highest predicted number.

Final Thoughts

Overall I was very impressed by Tensorflow.js and the accuracy I got considering that it was running on the browser. One other thing to mention is that you don't have to train your own model. You can use a pre-trained model for object detection, image caption, and more. The API is very simple to use and learn.