Skip to content

Model Loading

A model is a collection of artifacts that is created by the training process. In deep learning, running inference on a Model usually involves pre-processing and post-processing. DJL provides a ZooModel class, which makes it easy to combine data processing with the model.

This document will show you how to load a pre-trained model in various scenarios.

Using the ModelZoo API to load a Model

We recommend you use the ModelZoo API to load models.

The ModelZoo API provides a unified way to load models. The declarative nature of this API allows you to store model information inside a configuration file. This gives you great flexibility to test and deploy your model. See our reference project: DJL Spring Boot Starter.

Criteria class

You can use the Criteria class to narrow down your search condition and locate the model you want to load. Criteria class follows DJL Builder convention. The methods start with set are required fields, and opt for optional fields. You must call setType() method when creating a Criteria object:

Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class)
        .build();

The criteria accept the following optional information:

  • Engine: defines on which engine you want your model to be loaded
  • Device: defines on which device (GPU/CPU) you want your model to be loaded
  • Application: defines model application
  • Input/Output data type: defines desired input and output data type
  • artifact id: defines the artifact id of the model you want to load, you can use fully-qualified name that includes group id
  • group id: defines the group id of the pre-loaded ModelZoo that the model belongs to
  • ModelZoo: specifies a ModelZoo in which to search model
  • model urls: a comma delimited string defines at where the models are stored
  • Translator: defines custom data processing functionality to be used to ZooModel
  • Progress: specifies model loading progress
  • filters: defines search filters that must match the properties of the model
  • options: defines engine/model specific options to load the model
  • arguments: defines model specific arguments to customize the behavior of Translator

Note: If multiple models match the criteria you specified, the first one will be returned. The result is not deterministic.

Load model from the ModelZoo repository

The advantage of using the ModelZoo repository is it provides a way to manage models versions. DJL allows you to update your model in the repository without conflict with existing models. The model consumer can pick up new models without any code changes. DJL searches the classpath and locates the available ModelZoos in the system.

DJL provide several built-in ModelZoos:

You can create your own model zoo if needed, but we are still working on improving the tools to help create custom model zoo repositories.

Load models from the local file system

The following shows how to load a pre-trained model from a file path:

Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class) // defines input and output data type
        .optTranslator(ImageClassificationTranslator.builder().setSynsetArtifactName("synset.txt").build())
        .optModelPath(Paths.get("/var/models/my_resnet50")) // search models in specified path
        .optModelName("model/resnet50") // specify model file prefix
        .build();

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

DJL supports loading a pre-trained model from a local directory or an archive file.

Current supported archive formats

  • .zip
  • .tar
  • .tar.gz, .tgz, .tar.z

Customize modelName

By default, DJL will use the directory/file name of the URL as the model's modelName. DJL uses modelName as model file's prefix to load model file in the directory. We recommend naming the model file name to be the same as the directory or archive file.

If your model file located in a sub-folder of the model directory or has a different name, you can specify modelName by .optModelName() in criteria:

Criteria<Image, Classifications> criteria = Criteria.builder()
        .optModelName("traced_model/resnet18.pt") // specify model file prefix

You can also use the URL query string to tell DJL how to load model:

file:///var/models/resnet.zip?model_name=saved_model/resnet-18

Load model from a URL

DJL supports loading a model from a URL. Since a model consists multiple files, some of URL must be an archive file.

Current supported URL scheme:

  • file:// load a model from local directory or archive file
  • http(s):// load a model from an archive file from web server
  • jar:// load a model from an archive file in the class path
  • djl:// load a model from the model zoo
  • s3:// load a model from S3 bucket (requires djl aws extension)
  • hdfs:// load a model from HDFS file system (requires djl hadoop extension)
Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class) // defines input and output data type
        .optTranslator(ImageClassificationTranslator.builder().setSynsetArtifactName("synset.txt").build())
        .optModelUrls("https://resources.djl.ai/benchmark/squeezenet_v1.1.tar.gz") // search models in specified path
        .build();

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

You can customize the artifactId and modelName the same way as loading model from the local file system.

Load model from AWS S3 bucket

DJL supports loading a model from an S3 bucket using s3:// URL and the AWS plugin. See here for details.

Load model from Hadoop HDFS

DJL supports loading a model from a Hadoop HDFS file system using hdfs:// URL and the Hadoop plugin. See here for details.

Implement your own Repository

You may want to create additional model zoos using other protocols such as:

  • ftp://
  • sftp://
  • tftp://
  • rsync://
  • smb://
  • mvn://
  • jdbc://

DJL is highly extensible and our API allows you to create your own URL protocol handling by extending Repository class:

  • Create a class that implements RepositoryFactory interface make sure getSupportedScheme() returns URI schemes that you want to handle
  • Create a class that implements Repository interface.
  • DJL use ServiceLoader to automatically register your RepositoryFactory. You need add a file META-INF/services/ai.djl.repository.RepositoryFactory See java ServiceLoader for more detail.

You can refer to AWS S3 Repostory for an example.

Configure model zoo search path

DJL provides a way for developers to configure a system wide model search path by setting a ai.djl.repository.zoo.location system properties:

-Dai.djl.repository.zoo.location=https://djl-ai.s3.amazonaws.com/resnet.zip,s3://djl-misc/test/models,file:///myModels

The value can be comma delimited url string.

Debug model loading issues

You may run into ModelNotFoundException issue. In most cases, it's caused by the Criteria you specified doesn't match the desired model.

Here is a few tips you can use to help you debug model loading issue:

Enable debug log

See here for how to enable debug log

List models programmatically in your code

You can use ModelZoo.listModels() API to query available models.

List available models using DJL command line

Use the following command to list models in the DJL model zoo:

./gradlew :examples:listmodels

[INFO ] - CV.ACTION_RECOGNITION ai.djl.mxnet:action_recognition:0.0.1 {"backbone":"vgg16","dataset":"ucf101"}
[INFO ] - CV.ACTION_RECOGNITION ai.djl.mxnet:action_recognition:0.0.1 {"backbone":"inceptionv3","dataset":"ucf101"}
[INFO ] - CV.IMAGE_CLASSIFICATION ai.djl.zoo:resnet:0.0.1 {"layers":"50","flavor":"v1","dataset":"cifar10"}
[INFO ] - CV.IMAGE_CLASSIFICATION ai.djl.zoo:mlp:0.0.2 {"dataset":"mnist"}
[INFO ] - NLP.QUESTION_ANSWER ai.djl.mxnet:bertqa:0.0.1 {"backbone":"bert","dataset":"book_corpus_wiki_en_uncased"}

...

You can list models from your model folder with debug log:

./gradlew :examples:listmodels -Dai.djl.logging.level=debug -Dai.djl.repository.zoo.location=file:///mymodels