Skip to content

Run this notebook online:Binder

Load PyTorch model

In this tutorial, you learn how to load an existing PyTorch model and use it to run a prediction task.

We will run the inference in DJL way with example on the pytorch official website.


This tutorial requires the installation of Java Kernel. For more information on installing the Java Kernel, see the README.

// %mavenRepo snapshots

%maven ai.djl:api:0.22.1
%maven ai.djl.pytorch:pytorch-engine:0.22.1
%maven org.slf4j:slf4j-simple:1.7.32
import java.nio.file.*;
import java.awt.image.*;
import ai.djl.*;
import ai.djl.inference.*;
import ai.djl.modality.*;
import ai.djl.repository.zoo.*;
import ai.djl.translate.*;

Step 1: Prepare your model

This tutorial assumes that you have a TorchScript model. DJL only supports the TorchScript format for loading models from PyTorch, so other models will need to be converted. A TorchScript model includes the model structure and all of the parameters.

We will be using a pre-trained resnet18 model. First, use the DownloadUtils to download the model files and save them in the build/pytorch_models folder"", "build/pytorch_models/resnet18/", new ProgressBar());
Downloading: 100% |████████████████████████████████████████|

In order to do image classification, you will also need the synset.txt which stores the classification class labels. We will need the synset containing the Imagenet labels with which resnet18 was originally trained."", "build/pytorch_models/resnet18/synset.txt", new ProgressBar());
Downloading: 100% |████████████████████████████████████████| synset.txt

Step 2: Create a Translator

We will create a transformation pipeline which maps the transforms shown in the PyTorch example.

preprocess = transforms.Compose([
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

Then, we will use this pipeline to create the Translator

Translator<Image, Classifications> translator = ImageClassificationTranslator.builder()
        .addTransform(new Resize(256))
        .addTransform(new CenterCrop(224, 224))
        .addTransform(new ToTensor())
        .addTransform(new Normalize(
            new float[] {0.485f, 0.456f, 0.406f},
            new float[] {0.229f, 0.224f, 0.225f}))

Step 3: Load your model

Next, we add some search criteria to find the resnet18 model and load it. In this case, we need to tell Criteria where to locate the model by calling .optModelPath() API.

Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class)
        .optOption("mapLocation", "true") // this model requires mapLocation for GPU
        .optProgress(new ProgressBar()).build();

ZooModel model = criteria.loadModel();
Loading:     100% |████████████████████████████████████████|

[IJava-executor-0] INFO ai.djl.pytorch.engine.PtEngine - PyTorch graph executor optimizer is enabled, this may impact your inference latency and throughput. See:
[IJava-executor-0] INFO ai.djl.pytorch.engine.PtEngine - Number of inter-op threads is 1
[IJava-executor-0] INFO ai.djl.pytorch.engine.PtEngine - Number of intra-op threads is 2

Step 4: Load image for classification

We will use a sample dog image to run our prediction on.

var img = ImageFactory.getInstance().fromUrl("");