Skip to content

Latest commit

 

History

History
110 lines (85 loc) · 6.23 KB

README.md

File metadata and controls

110 lines (85 loc) · 6.23 KB
Raw

DJL - Apache MXNet model zoo

Introduction

The model zoo contains symbolic models from Apache MXNet (incubating) that can be used for inference and training. All the models in this model zoo contain pre-trained parameters for their specific datasets.

Documentation

The latest javadocs can be found on here.

You can also build the latest javadocs locally using the following command:

# for Linux/macOS:
./gradlew javadoc

# for Windows:
..\..\gradlew javadoc

The javadocs output is built in the build/doc/javadoc folder.

Installation

You can pull the MXNet engine from the central Maven repository by including the following dependency in you pom.xml file:

<dependency>
    <groupId>ai.djl.mxnet</groupId>
    <artifactId>mxnet-model-zoo</artifactId>
    <version>0.19.0</version>
</dependency>

Pre-trained models

The MXNet model zoo contains two major categories: Computer Vision (CV) and Natural Language Processing (NLP). All the models are grouped by task under these two categories as follows:

  • CV
    • Action Recognition
    • Image Classification
    • Object Detection
    • Pose Estimation
    • Semantic Segmentation/Instance Segmentation
  • NLP
    • Question and Answer

How to find a pre-trained model in model zoo

In a model zoo repository, there can be many pre-trained models that belong to the same model family. You can use the ModelZoo class to search for the model that you need. First, decide which model family you want to use. Then, define your key/values search criteria to narrow down the model you want. If there are multiple models that match your search criteria, the first model found is returned. ModelNotFoundException will be thrown if no matching model is found.

The following is an example of the criteria to find a Resnet50-v1 model that has been trained on the imagenet dataset:

Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class)
        .optArtifactId("resnet")
        .optFilter("layers", "50")
        .optFilter("flavor", "v1")
        .optFilter("dataset", "imagenet")
        .optDevice(device)
        .build();

ZooModel<Image, Classifications> model = criteria.loadModel();

List of search criteria for each model

The following table illustrates the possible search criteria for all models in the model zoo:

Category Application Model Family Criteria Possible values
CV Action Recognition ActionRecognition backbone vgg16, inceptionv3
dataset ucf101
Image Classification MLP dataset mnist
Resnet layers 18, 34, 50, 101, 152
flavor v1, v2, v1d
dataset imagenet, cifar10
Resnext layers 101, 150
flavor 32x4d, 64x4d
dataset imagenet
Senet layers 154
dataset imagenet
SeResnext layers 101, 150
flavor 32x4d, 64x4d
dataset imagenet
Instance Segmentation mask_rcnn backbone resnet18, resnet50, resnet101
flavor v1b, v1d
dataset coco
Object Detection SSD size 300, 512
backbone vgg16, mobilenet, resnet18, resnet50, resnet101, resnet152
flavor atrous, 1.0, v1, v2
dataset coco, voc
Pose Estimation SimplePose backbone resnet18, resnet50, resnet101, resnet152
flavor v1b, v1d
dataset imagenet
NLP Question and Answer BertQA backbone bert
dataset book_corpus_wiki_en_uncased

Note: Not all combinations in the above table are available. For more information, see the metadata.json file in the src/test/resources/mlrepo/model folder.

Contributor Guides and Documentation

How to add new models to the model zoo