Support string to IValue dict in stringMapFrom for pytorch

[jiyuanq]

Description

I’m planning to use djl mainly for serving pytorch models given it seems to be the most performant option available apart from implementing a service in c++. The input type of my model is quite complex so it’s not possible for me to use the Predictor api which feeds NDList to model. As a result I decided to work directly with the pytorch engine.

When I was researching how to construct an IValue that matches my model’s input type, I realized the current JNI layer only supports Map<String, NDArray>, but not Map<String, IValue>.

I think supporting Map<String, IValue> will make it much more flexible for anyone working with pytorch engine. In my case, for example, I need to create Map<String, Tuple<Tensor, Tensor>>, which is impossible right now.

If I’m understanding correctly it should not be a lot of effort to add this support given Tensor (NDArray) is itself an IValue base type

[frankfliu]

We do support IValue in PyTorch. But IValue is PyTorch specific, so didn't make it available at DJL API level. So if you want to use IValue, you have to write engine specific code.

See test code: https://github.com/deepjavalibrary/djl/blob/master/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java#L37

When you using IValue, you must be care to close all IValue you created, otherwise will cause memory.

[jiyuanq]

@frankfliu thanks for your reply! Just to clarify, I’m already using IValue, and I’m asking if we can enhance the IValue creation api.

In the existing IValue api, there’s a method IValue.stringMapFrom which takes a Map<String, PtNDArray>, but in my case I need to create an IValue from Map<String, IValue>

[frankfliu]

@jiyuanq No, we currently don't support it. I think PyTorch java api doesn't support it either. Feel free add add a PR to add Map<String, IValue> support if you need such feature.

[jiyuanq]

@frankfliu By “pytorch Java API” do you mean this? https://pytorch.org/javadoc/1.9.0/ it looks like it’s supported there, see “dictStringKeyFrom”

I think the c++ api supports generic dict of Dict<IValue, IValue>, so I’m guessing even when we pass in a Map<String, PtNDArray>, underlying the keys and values will be wrapped as IValues

[frankfliu]

@jiyuanq I see, now I remember, we dropped Dict(str, IValue) in favor of Dict(str, Tensor), C++ API support generic dict, but it's a bit hard to map in java, for simplicity, we only exposed Dict(str, Tensor)

We cannot have both:

    public static IValue stringMapFrom(Map<String, PtNDArray> map);
    public static IValue stringMapFrom(Map<String, IValue> map);

The following works, but a bit ugly:

    public static IValue stringTensorMapFrom(Map<String, PtNDArray> map);
    public static IValue stringIValueMapFrom(Map<String, IValue> map);

[jiyuanq]

@frankfliu I see. Thanks for the additional context! Is it possible to add stringIValueMapFrom as a new method?

I took a look at the code and it looks like it doesn't involve too much change. Just some changes to Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueFromStringMap in ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc plus corresponding java changes in IValue.java.

I was thinking about taking a stab myself but the build and release process seems quite complex, and the documentation about testing the JNI layer end to end seems rather sparse.

It would be great if someone familiar with the dev process can help add the support. What do you think?

[frankfliu]

@jiyuanq We will consider to priorities this feature.

[siddvenk]

This functionality is now included in DJL https://github.com/deepjavalibrary/djl/blob/master/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java#L313.