Loading Custom Torchscript Model in DJL

[tandeshao]

Hi, I would like to get the embeddings from running the vision transformer (https://huggingface.co/docs/transformers/v4.44.2/en/model_doc/vit) with DJL and i am wondering if it is possible?

Currently I have converted my vision transformer model in python to torchscript from this tutorial here: https://djl.ai/docs/pytorch/how_to_convert_your_model_to_torchscript.html.

However, i have trouble running the model.

Here is my rough code:

However, i seem to receive the following error:

May I know what does this error mean? Also is returning an NDList the correct way to retrieve the embeddings of the model?

[frankfliu]

@tandeshao
Your error most likely is caused by version mismatch between api and pytorch-engine. Can you share your pom.xml or build.gradle file?

We recommend you to use DJL BOM to ensure versions between different DJL components: https://docs.djl.ai/master/bom/index.html#djl-bill-of-materials-bom

[tandeshao]

Indeed, it seems like a POM issue. However, after changing my dependencies and running my code, i am met with this error:

No PyTorch native library matches your operating system: cpu-osx-x86_64:2.4.0

Here is my POM.xml

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example</groupId>
    <artifactId>testing_mvn</artifactId>
    <version>1.0-SNAPSHOT</version>
    <dependencyManagement>
        <dependencies>
            <dependency>
                <groupId>ai.djl</groupId>
                <artifactId>bom</artifactId>
                <version>0.30.0</version>
                <type>pom</type>
                <scope>import</scope>
            </dependency>
        </dependencies>
    </dependencyManagement>
    <dependencies>
        <dependency>
            <groupId>ai.djl</groupId>
            <artifactId>api</artifactId>
        </dependency>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-engine</artifactId>
        </dependency>
    </dependencies>

    <properties>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>

</project>

my logic:

my OS: macOS apple m2

[frankfliu]

PyTorch stopped intel based mac support. DJL stopped intel based mac since 0.30.0: https://docs.djl.ai/master/engines/pytorch/pytorch-engine/index.html#macos

You have to use 0.29.0 if you want to use old mac machine.

[tandeshao]

@frankfliu I see, i have updated my DJL version to 0.29.0, but the issue still persists. Do i need to add additional dependencies? Also FYI, here is my macbook specs:

I am using apple m2 chip

Here is my pom

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example</groupId>
    <artifactId>testing_mvn</artifactId>
    <version>1.0-SNAPSHOT</version>
    <dependencyManagement>
        <dependencies>
            <dependency>
                <groupId>ai.djl</groupId>
                <artifactId>bom</artifactId>
                <version>0.29.0</version>
                <type>pom</type>
                <scope>import</scope>
            </dependency>
        </dependencies>
    </dependencyManagement>
    <dependencies>
        <dependency>
            <groupId>ai.djl</groupId>
            <artifactId>api</artifactId>
        </dependency>
        <dependency>
            <groupId>ai.djl.pytorch</groupId>
            <artifactId>pytorch-engine</artifactId>
        </dependency>
    </dependencies>

    <properties>
        <maven.compiler.source>8</maven.compiler.source>
        <maven.compiler.target>8</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>

</project>

[frankfliu]

Your mac is M2, not x86_64, it's wired you got error: No PyTorch native library matches your operating system: cpu-osx-x86_64:2.4.0

Can you clone DJL repo, and run the following command:

cd djl
./gradlew debugEnv

[tandeshao]

@frankfliu After changing my Java SDK, i was able to load the torchscript model but encountered the error: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead

From my understanding, this is because apple's GPU does not support float64 operations. So i have instantiated the newPredictor without a device (which from my understanding defaults to CPU)

Running the code again results in a null pointer exception at this code here:

Specifically, block is null. Could you shed some light on how i may be able to solve this?

[tandeshao]

As reference, this is how i am saving the vision transformer as a torchscript file

[frankfliu]

@tandeshao

Yes, please avoid using MPS, it has many limitations.

It seems you didn't load your model yet. Please use Criteria API to load model: https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java#L163-L176

[frankfliu]

Another recommendation, please avoid NDList as output, better convert NDList to plain java object in your Translator. It's a bit tricky, you need fully understand NDArray lifecycle when you use NDList as output.

[tandeshao]

@frankfliu I see, are there any suggestions on how i could retrieve the model's embeddings from its last hidden state in djl? In python, i could do it like so:

[frankfliu]

google/vit-base-patch16-224 is an image-classification model, you should be able to use built-in ImageClassificationTranslatorFactory:

        Criteria<Image, Classifications> criteria =
                Criteria.builder()
                        .setTypes(Image.class, Classifications.class)
                        .optModelPath(path)
                        .optTranslatorFactory(new ImageClassificationTranslatorFactory())
                        .optArgument("flag", "GRAYSCALE")
                        .build();

[tandeshao]

I have tried changing the code to have the criteria logic as above. I ran to the following error:
Exception in thread "main" ai.djl.translate.TranslateException: java.io.FileNotFoundException: File not found: ./synset.txt

Below is my code:

public static void main(String[] args) throws ModelNotFoundException, MalformedModelException, IOException, TranslateException {
        List<NDList> embeddings = new ArrayList<>();

        // Logic to load images from url (this works)
        Image img = ImageFactory.getInstance().fromUrl("https://upload.wikimedia.org/wikipedia/en/a/a8/HELLOWORLD.jpg");
        float[] mean = new float[]{0.5f, 0.5f, 0.5f, 0.5f};
        float[] std = new float[]{0.5f, 0.5f, 0.5f, 0.5f};

        // Create a custom translator for the preprocessing and embedding extraction
        Criteria<Image, Classifications> criteria =
                Criteria.builder()
                        .setTypes(Image.class, Classifications.class)
                        .optModelPath(Paths.get("./vit.pt"))
                        .optTranslatorFactory(new ImageClassificationTranslatorFactory())
                        .optArgument("flag", "GRAYSCALE")
                        .build();

        try (ZooModel<Image, Classifications> model = criteria.loadModel();
             Predictor<Image, Classifications> predictor = model.newPredictor()) {
            System.out.println(predictor.predict(img));
        }
    }

[tandeshao]

@frankfliu After putting a random synset.txt file in my root repository, i believe i was able to successfully run my model. However, the pytorch engine encountered a runtime error:

RuntimeError: The size of tensor a (375) must match the size of tensor b (197) at non-singleton dimension 1
	at ai.djl.inference.Predictor.batchPredict(Predictor.java:195)
	at ai.djl.inference.Predictor.predict(Predictor.java:132)

Which i suspect is because the preprocessing step of the image is not done properly.

In python, the image processor is obtained via ViTImageProcessor from the transformer module. I am wondering if i am able to do the equivalent in djl? Or would i need to write the preprocessing logic myself?

If i am required to write myself, how can i do so? From the vision transformer paper, i know that images are scaled 384x384 and normalized with RGB channels with mean (0.5, 0.5, 0.5) and standard deviation (0.5, 0.5, 0.5). Is there any example i can follow to implement this?

[frankfliu]

@tandeshao
This is most likely related to pre-processing (which is Translator's job):

1. Your image is not grayscale, so you should not pass grayscle flag
2. You need to resize the image the image to 3x224x224
3. The model actually is a feature extraction model, you should use ImageFeatureExtractorFactory

You can always write your own Translator, but for cv tasks, DJL provides a rich set of built-in pre/post processing functionalities. You only need configure the argument in Criteria API.

        Criteria<Image, float[]> criteria =
                Criteria.builder()
                        .setTypes(Image.class, float[].class)
                        .optModelPath(modelFile)
                        .optEngine("PyTorch")
                        .optArgument("resize", true)
                        .optArgument("normalize", "0.5,0.5,0.5,0.5,0.5,0.5")
                        .optTranslatorFactory(new ImageFeatureExtractorFactory())
                        .build();

        try (ZooModel<Image, float[]> model = criteria.loadModel();
                Predictor<Image, float[]> predictor = model.newPredictor()) {
            float[] detection = predictor.predict(img);
        }

[tandeshao]

@frankfliu gotcha, after changing the code, the error changed to No model with the specified URI or the matching Input/Output type is found.

Here is the modified code based on what you mentioned above:

Perhaps there is something wrong with my model conversion? This is my python code that converts the vit model to a torchscript:

[frankfliu]

@tandeshao
You can find all supported image pre-processing configuration in the code: https://github.com/deepjavalibrary/djl/blob/master/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java#L155

[tandeshao]

@frankfliu After changing the output type to byte[], the code was able to run. However, the results of my output seems a little different (compared to when i ran it in pytorch). I see that in processOutput of ImageFeatureExtractor, the byte[] is generated from list.get(0).toByteArray(). Is this the equivalent of outputs.last_hidden_state[:, 0] in PyTorch?

Specifically here is the difference:
pytorch result:

java result:

[frankfliu]

Please use 0.31.0-SNAPSHOT version of DJL.

It should be float[], see: https://github.com/deepjavalibrary/djl/blob/master/api/src/main/java/ai/djl/modality/cv/translator/ImageFeatureExtractor.java#L38

[tandeshao]

@frankfliu as i will be using this code in production environment, i am not allowed to use any snapshot version of djl. Are there any workaround with existing djl versions?

[frankfliu]

1. DJL 0.31.0 is planed to release in Nov, you can test with SNAPSHOT for now and switch to release version once it released to maven central
2. You can always write your own Translator, just copy latest ImageFeatureExtractor

[tandeshao]

@frankfliu After running the model with my custom feature extractor, the result of the preprocessing step leads to a slight difference in the one run in pytorch.

Here is my criteria:

Criteria<Image, float[]> criteria =
                Criteria.builder()
                        .setTypes(Image.class, float[].class)
                        .optModelPath(Paths.get("./vit.pt"))
                        .optEngine("PyTorch")
                        .optArgument("height", 244)
                        .optArgument("width", 244)
                        .optArgument("resize", true)
                        .optArgument("normalize", "0.5,0.5,0.5,0.5,0.5,0.5")
                        .optTranslatorFactory(new CustomFeatureExtractorFactory())
                        .build();

The result of the preprocessed image in pytorch:

The result of the preprocessed image in djl:

Is this difference to be expected? Because it affects the result of my model.

Note: Even in the first preprocessing step (image resizing) the result of pytorch seems different from the one in djl:

Result of resizing in djl:

Result of resizing in pytorch:

[frankfliu]

@tandeshao
The different comes from image processing. First of all, you need to make sure you apply the same preprocessing in the code, e.g, if crop is applied, how the crop is performed, what resize algorithm is used.
However, it's hard to guarantee produce exactly the same result. In DJL, we use Java native image library (or OpenCV) to read the image and resize the image, while Python use PIL, this cause slightly difference in floating point result (usually should not effect result much).

[tandeshao]

@frankfliu I see, this makes sense. I have checked that indeed the difference between most operations are miniscule. Just to clarify, is this the correct way to setup width and height arguments in criteria?

Criteria<Image, float[]> criteria =
                Criteria.builder()
                        .setTypes(Image.class, float[].class)
                        .optModelPath(Paths.get("./vit.pt"))
                        .optEngine("PyTorch")
                        .optArgument("height", 244)
                        .optArgument("width", 244)
                        .optArgument("resize", true)
                        .optArgument("normalize", "0.5,0.5,0.5,0.5,0.5,0.5")
                        .optTranslatorFactory(new CustomFeatureExtractorFactory())
                        .build();

It seems like whenever i set this optional height and width, i receive a runtime error when doing prediction with the model:
RuntimeError: The size of tensor a (226) must match the size of tensor b (197) at non-singleton dimension 1

However, the model predicts fine when the two optArguments are removed.

[frankfliu]

the size should be 224 not 244.

If you resize the image to 224x224 before calling DJL or Python, the result should be identical.