The Future of Machine Learning in the Browser with TensorFlow.js

Up until now, I did most of my machine learning work in Python. With the recent release of TensorFlow.js – TensorFlow for JavaScript, I decided to spend the last couple of months trying out machine learning in the browser.

I am well aware of an intensive relationship people have with JavaScript. But let’s be honest – JavaScript entering the machine learning world was inevitable. Jeff Atwood has a famous quote “Any application that can be written in JavaScript, will eventually be written in JavaScript”.

I wanted to decide for myself whether doing machine learning in a browser is just a gimmick or something worth to invest more time in the near future. Here’s what I’ve learned.

This is my guest post from Intermix.io blog.

TensorFlow.js: Easy to setup, interactive and real-time

Visiting a website with a TensorFlow.js model on it is like visiting any regular website. Users don’t have to go through any particular setup – it just works. Having to run a Python-based machine learning algorithm on the user’s device is a whole different story. To do that, we need to compile the libraries with the application code. But usually, we don’t bother and delegate machine learning to the server completely.

Given my background in web development, building and training machine learning models in the browser feel magical. The available input is richer with an easy access to a webcam and a microphone. It’s also possible to re-use the same code on smartphones, with access to gyroscope and GPS sensors, to name a few.

This, combined with the real-time nature of running the code on the client side, enables highly-interactive gaming experiences. Google Brain Team released a humorous Pac-man game version where you control your character using your head. You train a machine learning model to associate your head movements with keyboard controls.

Image 1: Pac-man, controllable with head movements.

TensorFlow.js also allows debugging machine learning models more easily. With TensorFlow in Python, there was always a bit of work involved in understanding what my model was doing. Countless rows of print statements and TensorBoard certainly help but cannot compete with having a browser window to have my results and errors rendered in. There are numerous education examples that showcase how to debug models even on an individual neuron level.

Image 2: A Neural Network Playground Daniel Smilkov and Shan Carter.

Doing machine learning in the browser enables a different kind of developer experience. Simple setup, easily accessible sensors, real-time nature of the client side along with highly interactive tooling have a great impact on my happiness as a developer.

User’s data can stay on the user’s machine

In the age of Cambridge Analytica scandals, people are hesitant to give up their data to companies. Does that mean that they should not take advantage of the recent machine learning advancements?

An obvious, yet very important aspect of TensorFlow.js is that it can run directly on the user’s device – without ever needing to communicate with any server. Sending voice recordings to servers in Detroit just to do “Ok Google” style hot-word activation is dangerous. We can now do it on the device itself. On the flip side, companies might not even want to touch your data themselves since storing user’s data poses certain liabilities.

It also allows us to build browser extensions that adapt to your browsing habits while keeping them private. Here’s an experimental Chrome extension that recognizes images while you browse the internet.

Performance on par with the Python version

I’ve always been conservative about JavaScript’s performance and did not expect miracles. Performance is a key factor when running machine learning models and that’s why we all rely on GPUs and TPUs.

To train a model in the browser, TensorFlow.js needs to access the GPU through WebGL. Large models might train up to 10-15x slower when compared to TensorFlow Python with AVX. The good part is that most likely you aren’t doing intensive training on user’s devices. And running the trained model is only 1.5-2x slower. That scales so much better than sending an image from a browser to the servers, doing inference there and sending the results back to the browser.

Being slower is not a final verdict. Recently released (not yet production-ready) version of TensorFlow.js for Node.js is on par with the Python CUDA. That’s right – JavaScript TensorFlow is exactly as fast as Python version on the same top-of-the-line hardware.

Image 3: Performance comparison from a recent talk at Google I/O ’18.

In other words – performance of TensorFlow.js might not be a valid concern in a few months.

Wrap up

At this point, machine learning in the browser is in its early stages. There were Brain.js and Synaptic, but none of them had the support of the entire Google Brain team and were nowhere near full feature parity with TensorFlow Python. True, running TensorFlow.js on WebGL is slower, Node.js version is not yet production-ready and some features are still missing. Yet, running machine learning models both on the client and server might be the beginning of new era of ML-enabled web apps.