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.

Preparation

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

// %mavenRepo snapshots https://oss.sonatype.org/content/repositories/snapshots/

%maven ai.djl:api:0.8.0
%maven ai.djl.pytorch:pytorch-engine:0.8.0
%maven org.slf4j:slf4j-api:1.7.26
%maven org.slf4j:slf4j-simple:1.7.26
%maven net.java.dev.jna:jna:5.3.0

// See https://github.com/awslabs/djl/blob/master/pytorch/pytorch-engine/README.md
// for more PyTorch library selection options
%maven ai.djl.pytorch:pytorch-native-auto:1.6.0
import java.awt.image.*;
import ai.djl.*;
import ai.djl.inference.*;
import ai.djl.modality.*;
import ai.djl.modality.cv.*;
import ai.djl.modality.cv.util.*;
import ai.djl.modality.cv.transform.*;
import ai.djl.modality.cv.translator.*;
import ai.djl.repository.zoo.*;
import ai.djl.translate.*;
import ai.djl.training.util.*;

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

DownloadUtils.download("https://djl-ai.s3.amazonaws.com/mlrepo/model/cv/image_classification/ai/djl/pytorch/resnet/0.0.1/traced_resnet18.pt.gz", "build/pytorch_models/resnet18/resnet18.pt", new ProgressBar());
Downloading: 100% |████████████████████████████████████████| resnet18.pt

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.

DownloadUtils.download("https://djl-ai.s3.amazonaws.com/mlrepo/model/cv/image_classification/ai/djl/pytorch/synset.txt", "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.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    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

Pipeline pipeline = new Pipeline();
pipeline.add(new Resize(256))
        .add(new CenterCrop(224, 224))
        .add(new ToTensor())
        .add(new Normalize(
            new float[] {0.485f, 0.456f, 0.406f},
            new float[] {0.229f, 0.224f, 0.225f}));

Translator<Image, Classifications> translator = ImageClassificationTranslator.builder()
                .setPipeline(pipeline)
                .optApplySoftmax(true)
                .build();

Step 3: Load your model

Next, we will set the model zoo location to the build/pytorch_models directory we saved the model to. You can also create your own Repository to avoid manually managing files.

Next, we add some search criteria to find the resnet18 model and load it.

// Search for models in the build/pytorch_models folder
System.setProperty("ai.djl.repository.zoo.location", "build/pytorch_models/resnet18");

Criteria<Image, Classifications> criteria = Criteria.builder()
        .setTypes(Image.class, Classifications.class)
         // only search the model in local directory
         // "ai.djl.localmodelzoo:{name of the model}"
        .optArtifactId("ai.djl.localmodelzoo:resnet18")
        .optTranslator(translator)
        .optProgress(new ProgressBar()).build();

ZooModel model = ModelZoo.loadModel(criteria);
[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

Loading:     100% |████████████████████████████████████████|

Step 4: Load image for classification

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

var img = ImageFactory.getInstance().fromUrl("https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg");
img.getWrappedImage()