Deep Learning in Java
Keerthan Vasist medium.com
Java has been one of the most popular programming languages in enterprise for a long time and has a massive ecosystem of libraries, and frameworks, and a large community of developers. However, there…
https://miro.medium.com/v2/resize:fill:256:256/1*VzTUkfeGymHP4Bvav-T-lA.png
Medium
Java has been one of the most popular programming languages in enterprise for a long time and has a massive ecosystem of libraries, and frameworks, and a large community of developers. However, there are very limited options offered in Java for deep learning applications. Currently, most deep learning models are written and trained in Python. This presents an additional barrier to entry for Java developers who want to enter this area, as they have to learn both a new programming language and the complex domain of deep learning.

In order to lower the barrier of entry into deep learning for Java developers, AWS built Deep Java Library (DJL), an open source deep learning framework in Java to bridge the gap for Java developers by supporting any deep learning engine, such as Apache MXNet, PyTorch, or TensorFlow to run training and inference natively in Java. It also contains a powerful ModelZoo design that allows you to manage trained models and load them in a single line of code. The built-in ModelZoo currently supports more than 70 pre-trained and ready to use models from GluonCV, HuggingFace, TorchHub and Keras. If you’ve been a Java developer and are interested in exploring deep learning, Deep Java Library (DJL) is a great place to start.

In this tutorial, we walk through an example demonstrating the training capabilities of DJL by training a simple model on the popular MNIST dataset.

What is Deep Learning?

Machine learning is the process of letting the computer learn the specifications of the given task from data through the use of various statistical techniques. This ability to learn the features of a task allows computers to perform complex tasks, such as detecting an object in an image, that were generally considered to be beyond the scope of computers because of the difficulty of providing exact specifications for every possible case.

Deep learning is a branch of machine learning based on artificial neural networks. An artificial neural network is a programming paradigm inspired by the human brain that helps the computer learn and perform tasks based on observational data. Deep learning is a collection of powerful techniques that can be leveraged to help train large artificial neural networks to perform complex tasks. Deep learning techniques have proven to be very effective at solving complex tasks like object detection, action recognition, machine translation, natural language understanding among others.

Train MNIST with DJL

Setting up the project

You can use the following configuration in your gradle project to import the required dependencies. In this example, we use the api package which contains the core API of the DJL project, and the basicdataset package which contains some of the basic datasets in DJL. Since we are training with the MXNet engine, we are also going to import the mxnet-engine package, and the mxnet-native-auto package.

plugins {

id 'java'

}

repositories {

jcenter()

}

dependencies {

implementation "ai.djl:api:0.8.0"

implementation "ai.djl:basicdataset:0.8.0"

// MXNet

runtimeOnly "ai.djl.mxnet:mxnet-engine:0.8.0"

runtimeOnly "ai.djl.mxnet:mxnet-native-auto:1.7.0-backport"

}

NDArray and NDManager

NDArray is the core data structure for all mathematical computations in DJL. An NDArray represents a multidimensional, fixed-size homogeneous array. NDArray behaves similar to the python program numpy.

NDManager is the manager of NDArray. NDManager manages the lifecycle of NDArray, and is an essential part of memory management in DJL. Every NDArray created by an instance of NDManager will be closed once that NDManager is closed. NDManager, and NDArray both extend AutoCloseable. To know and understand the usage of NDArray and NDManager better, see this blog post.

Model

In DJL, training and inference start with a Model. In this article, we concentrate on the training process. To start the training process we create a new instance of the Model class. The Model class also extends AutoCloseable. So, it is created with a try-with-resources.

try (Model model = Model.newInstance()) {

...

// training process takes place here

...

}

Preparing the data

The MNIST database (Modified National Institute of Standards and Technology database) is a large database of handwritten digits that is commonly used for training various image processing systems. The MNIST dataset is readily available in DJL. The shape of an individual image from the MNIST dataset in DJL is (28, 28). If you wish to train a model on your own dataset, you can do that by adding your own dataset by following the instructions here.

In order to train your model, you first need to load the dataset.

int batchSize = 32;

Mnist trainingDataset = Mnist.builder()

.optUsage(Usage.TRAIN)

.setSampling(batchSize, true)

.build();

Mnist validationDataset = Mnist.builder()

.optUsage(Usage.TEST)

.setSampling(batchSize, true)

.build();

This code creates training and validation datasets. The dataset is also configured to sample the dataset randomly. There are further configurations made on the dataset like applying transforms on the images, or limiting the size of the dataset.

Building the model (Constructing the Block)

Once the data is ready, you need to construct the neural network that you want to train. In DJL, a neural network is represented by a Block. A Block is a composable function that forms a neural network. They can represent single operation, parts of a neural network, and even the whole neural network. A Block can have parameters and child blocks. During the training process, the parameters are updated, and the child blocks are also trained. This recursively updates the parameters of all its children as well.

When building these block functions, the easiest way is to use composition. Blocks can be built by combining other blocks. We refer to the containing block as the parent and the sub-blocks as the children.

We provide several helpers to make it easy to build common block composition structures. SequentialBlock is a container block whose children form a chain of blocks where each child block feeds its output to the next child block in a sequence. ParallelBlock is a container block whose children are executed in parallel, and the output of the blocks are combined according to a combining function specified. LambdaBlock is a block with an operating function that has to be specified by the user.

We are going to build a simple MLP (Multi-layer Perceptron). A multilayer perceptron (MLP) is a feedforward artificial neural network that generates a set of outputs from a set of inputs. An MLP is characterized by several layers of input nodes connected as a directed graph between the input and output layers. It can be constructed by using a number of LinearBlock within a SequentialBlock.

int input = 768;

int output = 10;

int[] hidden = new int[] {128, 64};

SequentialBlock sequentialBlock = new SequentialBlock();

sequentialBlock.add(Blocks.batchFlattenBlock(input));

for (int hiddenSize : hidden) {

sequentialBlock.add(Linear.builder().setUnits(hiddenSize).build());

sequentialBlock.add(activation);

}

sequentialBlock.add(Linear.builder().setUnits(output).build());

DJL also provides a pre-constructed Mlp block that we can use directly.

Block block = new Mlp(

Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH,

Mnist.NUM_CLASSES,

new int[] {128, 64});

Training

Now that you have created a new instance of the model, prepared the dataset, and constructed a block, you are ready to start training. In deep learning, training involves the following steps:

Initialization: This step initializes the blocks, and creates its corresponding parameters according to the specified Initializer schemes.

Forward: This step executes the computations represented by the Block, and generates the output.

Loss computation: In this step, you compute the loss by applying the specified Loss function to the output and label provided.

Backward: During this step, you use the loss, and back-propagate the gradients along the neural network.

Step: During this step, you update the values of the parameters of the block based on the specified Optimizer.

However, DJL abstracts all of these steps through the Trainer. The Trainer can be created by specifying the training configurations such as the Initializer, Loss, and Optimizer. These configurations and more can be set by using the TrainingConfig. Some of the other configurations that can be set are:

Device - the devices on which the training must occur

TrainingListeners - listeners that listen for various stages during the training process and execute specific functions like logging and evaluation. Users can implement custom TrainingListener as necessary.

DefaultTrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())

.addEvaluator(new Accuracy())

.optDevices(Device.getDevices(arguments.getMaxGpus()))

.addTrainingListeners(TrainingListener.Defaults.logging(arguments.getOutputDir()));

try (Trainer trainer = model.newTrainer(config)){

// training happens here

}

Once a trainer is created, it has to be initialized with the Shape of the inputs. Then, you call the fit() method to start training. The fit() method trains the model on the dataset for a specified number of epochs, runs validation, and saves the model in the specified directory in the filesystem.

/*

* MNIST is 28x28 grayscale image and pre processed into 28 * 28 NDArray.

* 1st axis is batch axis, we can use 1 for initialization.

*/

Shape inputShape = new Shape(1, Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH);

int numEpoch = 5;

String outputDir = "/build/model";// initialize trainer with proper input shape

trainer.initialize(inputShape);TrainingUtils.fit(trainer, numEpoch, trainingSet, validateSet, outputDir, "mlp");

That’s it. Congratulations! You have trained your first deep learning model using DJL! You can monitor the training process on the console, or however the listeners are implemented. If you use the default listeners, your output should be similar to the following.

[INFO ] - Downloading libmxnet.dylib ...

[INFO ] - Training on: cpu().

[INFO ] - Load MXNet Engine Version 1.7.0 in 0.131 ms.

Training: 100% |████████████████████████████████████████| Accuracy: 0.93, SoftmaxCrossEntropyLoss: 0.24, speed: 1235.20 items/sec

Validating: 100% |████████████████████████████████████████|

[INFO ] - Epoch 1 finished.

[INFO ] - Train: Accuracy: 0.93, SoftmaxCrossEntropyLoss: 0.24

[INFO ] - Validate: Accuracy: 0.95, SoftmaxCrossEntropyLoss: 0.14

Training: 100% |████████████████████████████████████████| Accuracy: 0.97, SoftmaxCrossEntropyLoss: 0.10, speed: 2851.06 items/sec

Validating: 100% |████████████████████████████████████████|

[INFO ] - Epoch 2 finished.NG [1m 41s]

[INFO ] - Train: Accuracy: 0.97, SoftmaxCrossEntropyLoss: 0.10

[INFO ] - Validate: Accuracy: 0.97, SoftmaxCrossEntropyLoss: 0.09

[INFO ] - train P50: 12.756 ms, P90: 21.044 ms

[INFO ] - forward P50: 0.375 ms, P90: 0.607 ms

[INFO ] - training-metrics P50: 0.021 ms, P90: 0.034 ms

[INFO ] - backward P50: 0.608 ms, P90: 0.973 ms

[INFO ] - step P50: 0.543 ms, P90: 0.869 ms

[INFO ] - epoch P50: 35.989 s, P90: 35.989 s

Once training is complete, we can use the trained model to run inference to get the predictions. You can follow the Inference with your model jupyter notebook to run inference on a saved model. You can also run the complete code directly by following the instructions in Train Handwritten Digit Recognition using Multilayer Perceptron (MLP) model.

Summary