Implement Object Detection with PyTorch in Java in 5 minutes

with DJL, an Engine-agnostic Deep Learning Library

PyTorch is one of most popular Machine Learning (ML) frameworks. Its easy to use Pythonic interface and productivity enhancing features, such as data parallelism and dynamic computation graphs, have simplified the development of ML models for applications in Computer vision (CV) and Natural Language processing (NLP).

While there are a variety of options to serve PyTorch models in production, there are only a few options to deploy PyTorch models natively in Java. Previously, users could write a Java wrapper around the PyTorch C++ API or use the experimental PyTorch Java bindings . Amazon’s Deep Java Library (DJL) now offers the PyTorch and Java community a simpler option with its easy to use high-level APIs. By abstracting the complexity involved in ML and bundling tedious data processing routines, DJL simplifies running predictions with a PyTorch model down to a few lines of code. In our two-part blog series, we demonstrate how users can run inference with PyTorch. First with pre-trained PyTorch models, then with user-generated PyTorch models.

In the first of this two-part blog, we showcase how you can implement object detection with a pre-trained PyTorch model in less than 5 minutes. So, let’s get started.

Setup

To get started with DJL, create a project using the command line or IDE of your choice. Add the following code to your build.gradle config file. We use the PyTorch 1.4.0 C++ distribution in this example.

Object Detection with a pre-trained PyTorch model

In this section, we implement Object Detection with a pretrained PyTorch traced model from NVIDIA . This task identifies objects embedded (dog, bike, truck) in the image below :

With DJL, you can run inference in just a few lines with the following code block:

Running the PyTorch code yields the following output. The output lists the objects detected along with their relative probability.

You can then use our Image visualization function to define bounding boxes with this image.

What is DJL?

After completing this tutorial, you may be wondering what DJL is. Deep Java Library (DJL) is an open source library to build and deploy deep learning in Java. This project launched in December 2019, and is widely used among teams at Amazon. This effort was inspired by other DL frameworks, but was developed from the ground up to better suit Java development practices. DJL is framework agnostic, with support for Apache MXNet, PyTorch, TensorFlow 2.x (Experimental) and fastText (Experimental).

PyTorch Support coverage

We built the DJL PyTorch Engine with the PyTorch (1.4.0) C++ API, which allows us to run operators and inference code. All models that are in the TorchScript format can be imported and run on DJL. To accomplish this, we wrote a customized JNI layer that interacts between C++ and Java. Currently, DJL covers more than 60 PyTorch operators. DJL itself has built-in features such as:

• NDManager: Efficient garbage collection system
• Multi-threading support (Experimental)
• ModelZoo: The PyTorch model zoo supports Image Classification and Object Detection models
• and lot more…

Supported Platforms:

• Mac: CPU
• Linux: CPU and GPU (CUDA 9.2 and CUDA 10.1)
• Windows: CPU and GPU (CUDA 9.2 and CUDA 10.1)

To learn more, check out our website , Github repository and Slack channel. Stay tuned for our next blog post where we will demonstrate how you can run inference with your own PyTorch models in Java.