Converting Essentia TensorFlow Models

Our parent library, Essentia, has a few machine learning models for music and audio available. These models are available here. Often, we can convert these TensorFlow models to be used together with Essentia.js and TensorFlow.js inside the browser. Here's how:

1. Get the model

First we need to download the model that we want to convert from https://essentia.upf.edu/models.html. Download the model's .pb file (which contains the model's weights) and its corresponding .json file (which contains useful metadata), and place them into a new folder.

mkdir model-conversion-folder
cd model-conversion-folder

2. Setup

Then we need to create an isolated python environment. There are two options:

a. Virtual environment

Inside the folder we created in step 1:

python3 -m venv .venv
source .venv/bin/activate
pip install tensorflowjs[wizard]

b. Docker

Using this as an example Dockerfile inside the folder we created in step 1:

FROM python:3.6
WORKDIR /model-conversion
RUN pip3 install --upgrade pip
RUN pip3 install tensorflowjs[wizard]

Then run these on your terminal:

docker build --tag essentiajs-model-conversion .
docker run -it -v `pwd`:/model-conversion/ essentiajs-model-conversion:latest bash

3. Convert

Lastly, we will use the tensorflowjs_converter tool that was installed with the tensorflowjs[wizard] package to convert the .pb model file into a .bin weights file and a corresponding .json file that TensorFlow.js can load. We will use a small version of the CREPE pitch detection model as an example here.

At this point we should have downloaded from https://essentia.upf.edu/models.html the crepe-small-1.pb and crepe-small-1.json files. The crepe-small subdirectory inside our model-conversion-folder was created to contain all of this model's files.

Example conversion command:

tensorflowjs_converter \
--input_format=tf_frozen_model \
--output_format=tfjs_graph_model \
--output_node_names='model/classifier/Sigmoid' \
./crepe-small/crepe-small-1.pb \
./crepe-small/tfjs

The value for the --output_node_names option ('model/classifier/Sigmoid' in this example) can be found in the .json file we downloaded (crepe-small-1.json), inside the schema.outputs[0].name property (where outputs[0] is the final layer and only available output in this particular model).