Pre-trained TensorFlow.js models

This repository hosts a set of pre-trained models that have been ported to TensorFlow.js.

The models are hosted on NPM and unpkg so they can be used in any project out of the box. They can be used directly or used in a transfer learning setting with TensorFlow.js.

To find out about APIs for models, look at the README in each of the respective directories. In general, we try to hide tensors so the API can be used by non-machine learning experts.

For those interested in contributing a model, please file a GitHub issue on tfjs to gauge interest. We are trying to add models that complement the existing set of models and can be used as building blocks in other apps.

Models

Image

• MobileNet - Classify images with labels from the ImageNet database.
  • npm i @tensorflow-models/mobilenet
• PoseNet - Realtime pose detection. Blog post here.
  • npm i @tensorflow-models/posenet
• Coco SSD - Object detection based on the TensorFlow object detection API.
  • npm i @tensorflow-models/coco-ssd

Audio

• Speech Commands - Classify 1 second audio snippets from the speech commands dataset.
  • npm i @tensorflow-models/speech-commands

Text

• Universal Sentence Encoder - A model that encodes English text into 512-dimensional embeddings.
• Text Toxicity - Detects and classifies toxic content such as threats, insults, and obscenities in English text inputs.

General utilities

• KNN Classifier - Create a custom k-nearest neighbors classifier. Can be used for transfer learning.
  • npm i @tensorflow-models/knn-classifier

Development

You can run the unit tests for any of the models by running the following inside a directory:

yarn test

New models should have a test NPM script.

To run all of the tests, you can run the following command from the root of this repo:

yarn presubmit