Pytorch model in Deep Java Library | by Shantanu Bhattacharyya | Towards Data Science

PyTorch is a fast growing and very popular open source Machine Learning framework. Its imperative design combined with “numpy” like workflow makes it a compelling first choice for beginners and professionals alike. However, serving these model in production is not straightforward and things are particularly difficult if the goal is to serve them natively in Java.
Amazon’s Deep Java Library (DJL) aims to solve this particular pain point by providing high level APIs that can run inference on PyTorch models with very little code. My recent test drive with DJL tells me that it could be a very powerful tool but the existing example set and community guidance (aka stackoverflow help :) ) could be a little intimidating for new folks, specially those who come from python background and are unfamiliar with the java style. This simple demo, hopefully, makes things easier for them. All the scripts are also available at my git repo here.
Let’s start by creating a simple linear regression model with PyTorch.
import torchimport torch.nn as nnX = torch.tensor([[1],[2],[4],[7],[9]], dtype = torch.float32)Y = torch.tensor([[2],[4],[8],[14],[18]], dtype = torch.float32)X_test = torch.tensor([[5]], dtype=torch.float32)n_sample, n_features = X.shapemodel = nn.Linear(n_features, n_features)learn_rate = 0.01n_epochs = 500loss = nn.MSELoss()optimizer = torch.optim.SGD(lr = learn_rate, params=model.parameters())for i in range(0,n_epochs):    y_pred = model(X)    ls = loss(y_pred,Y)    ls.backward()    optimizer.step()    optimizer.zero_grad()    [w,b] = model.parameters()    traced_cell = torch.jit.trace(model, X_test)    #print(traced_cell(X_test).item())    print(f"{ls.item():0.3f}, {w[0][0].item():0.3f {model(X_test).item():0.3f}")traced_cell.save('./model1.zip')
The model (model1) learns to take an input n and predict 2 * n.
There are 2 lines in the script above that are usually not present in PyTorch scripts.
1. traced_cell = torch.jit.trace(model, X_test) 2. traced_cell.save('./model1.zip')
These two lines are required to create a TorchScript model (serialized pytorch model) that can be used in a language agnostic fashion in high performance environments. We will now use this model to make inference using DJL. I have used a Jupyter notebook running Java kernel for this demo.
1 : Get all the required maven dependencies. DJL supports pytorch models out of the box. SLF4J is for logging purposes.
%maven ai.djl.pytorch:pytorch-engine:0.8.0%maven ai.djl.pytorch:pytorch-native-auto:1.6.0%maven org.slf4j:slf4j-simple:1.7.26
If you get errors that say “unknown resolver null”, you can ignore them for now.
2 : Import the required modules
import ai.djl.*;import ai.djl.inference.*;import ai.djl.ndarray.*;import ai.djl.translate.*;import java.nio.file.*;
It is worth highlighting the ndarray import here. N-Dimensional array is a popular data structure in many languages and numpy (ndarray library in python) is pretty much ubiquitous in the field of deep learning or pretty much any application involving numerical computation.
Java doesn’t have a similar ndarray implementation. DJL provides the ndarray functionality through the above import. This is one of the key factors that allows DJL to work with PyTorch models.
3. Declare path to model and model name. This assumes you are using a locally generated model. DJL provides some standard models too that can be used directly.
Path modelDir = Paths.get("Path_To_Your_Model_Folder/");Model model = Model.newInstance("model name");model.load(modelDir);model
Current DJL documentation says models can be .zip , .tar, .tar.gz, .tgz or .tar.z
If the model is correctly loaded, you should see an output like ai.djl.pytorch.engine.PtModel@4cefa68f
4. Translator to handle the pre-processing and post-processing steps.
Translator<Float, Float> translator = new Translator<Float, Float>(){    @Override    public NDList processInput(TranslatorContext ctx, Float input) {        NDManager manager = ctx.getNDManager();        NDArray array = manager.create(new float[] {input});        return new NDList (array);    }        @Override    public Float processOutput(TranslatorContext ctx, NDList list) {        NDArray temp_arr = list.get(0);        return temp_arr.getFloat();    }        @Override    public Batchifier getBatchifier() {        // The Batchifier describes how to combine a batch together        // Stacking, the most common batchifier, takes N [X1, X2, ...] arrays to a single [N, X1, X2, ...] array        return Batchifier.STACK;    }};
This is the step that confuses most people. The translator defines the signature for the input and output data types <Input type, Output type>. For our example they both happen to be a float. Make sure to use Float, not float, as a reference is explicitly expected. This translation step is not optional.
Further, if you look at the DJL doc for PreProcessor , you will find that NDList is the expected type for the preprocessing output that gets fed to the model.
So, even if your input was hypothetically already an NDList type, you still need to do this step.
The case for post processing is similar. It consumes the NDList and outputs the type described by your Translator.
All the three overides above are mandatory as well. There are of’course pre-implemented translators provided by DJL and I recommend you to explore them, particularly the image processing ones.
5. Making Prediction
Congratulations, you have done the hard work and are just a couple of lines of code away from making the inference. Here they are
Predictor<Float, Float> predictor = model.newPredictor(translator);predictor.predict(Your_Test_Float)
For my test case, I did predictor.predict(2.9f) and it output 5.799781. Ahh, pure Joy !
So we saw how we can create a PyTorch model and use that model for inference in a Java environment. This was a very simplistic use case though. For more advanced use cases, making the translator do the preprocessing and postprocessing correctly could require a few iterations. Further, the number of inbuilt modules (like the image processing ones) are currently few in number.
I am confident though, that DJL will evolve rapidly and could be the perfect library to utilise PyTorch models in Java (or C++) environment giving the perfect combination of development speed and high performance.