What java API can I use to extract the names from a PyTorch model file?

Mutant Bob 0
Viewed 27

I have been given a pytorch model file, and some object detection results. The object detection results give number to identify what kind of object it detected, but I want the names from the model file.

Some python code I found looks like this

model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, names, pt = model.stride, model.names, model.pt

I am pretty sure I need to get the names array, but I'm working in java, not python. I looked into the ai.djl.pytorch.engine.PtModel, but could not spot anything that looks like a mapping from numbers to names.

It even looks like DeepJavaLibrary can't support loading a plain .pt file:

String fname = "/tmp/yolov5s.pt"; 
{
    PtEngine engine = (PtEngine) Engine.getEngine("PyTorch");
    Model model = engine.newModel("bacon", null);
    model.load(new File(fname).toPath());
    Block block = model.getBlock();
    System.out.println(block);
}

Exception in thread "main" ai.djl.engine.EngineException: PytorchStreamReader failed locating file constants.pkl: file not found
    at ai.djl.pytorch.jni.PyTorchLibrary.moduleLoad(Native Method)
    at ai.djl.pytorch.jni.JniUtils.loadModule(JniUtils.java:1550)
    at ai.djl.pytorch.engine.PtModel.load(PtModel.java:90)
    at ai.djl.Model.load(Model.java:110)
    at project.pictureServer.PyTorchFile.main(PyTorchFile.java:37)

What is the proper way to map from object/class numbers to names using Java and a PyTorch model file?

java pytorch
1 Answers

Q: What is the proper way to map from a pyTorch model file to Java objects?

A: I believe a pyTorch model file is just a pickle serialization of Python objects. One alternative might be PythonPickle.

Your challenge is to read a PyTorch (.pt?) model file in Java.

  • One alternative is to "reverse engineer" those parts of the format you're interested in, and write your own ".pt decoder".

    It sounds like you've already started down this path, and had some success. If this works for you - great!

  • Since a ".pt" file is just "pickled Python", I suggested trying PythonPickle. It's very "lightweight", and it looked like it might do everything you're looking for - and more.

  • Yet another alternative might be DeepJavaLibrary, a set of Java-language bindings for pyTorch.

  • Who knows - you might even want to write a little Python script to a) read a .pt file, b) write it to JSON.

Anyway - please let us know what you decide.

paulsm40
Related