Problem
There are 2 phases to applying Deep Learning to a ML problem, the first phase is where a neural network is created and trained using training data to generate pre-trained model and In the second phase, this pre-trained model is put to work by running inference(forward-pass) on new data in the customer’s application in production. Model Creation and Training is typically performed by Data Scientists who prefer using Python as a primary language which provides rich set of libraries(numpy, pandas, pillow) etc., to setup the training pipeline. MXNet already has very good support for Python to quickly prototype and develop models.
Inference on the other hand is run and managed by Software Engineers in a production eco-system which is built with tools and frameworks that use Java/Scala as a primary Language.
Inference on a trained model has two different use-cases:
- Real time or Online Inference - tasks that require immediate feedback, such as fraud detection
- Batch or Offline Inference - tasks that don't require immediate feedback, these are use-cases where you have massive amounts of data and want to run Inference or pre-compute inference results
Batch Inference is performed on big data platforms such as Spark using Scala or Java while Real time Inference is typically performed and deployed on popular web frameworks such as Tomcat, Netty, Jetty, etc. which use Java.
With this project, we want to build Java APIs that are easy to use for Inference and lowers the entry barrier is consuming MXNet for production use-cases.
Goals
As a user, I’d like to have a Java Inference API that allows me to use deep learning models from my existing Java application
As a user, I’d like for the new Java inference API to be idiomatic and easy to use so that I can quickly learn to deploy models.
As a user, I’d like for the new Java Inference API to introduce as few dependencies as possible so that it’s easy to add into my existing environment.
As a user, I’d like for the new Java API to have full test coverage so that I can be confident in it’s stability.
As a user, I’d like for the Java API to be as performant as possible with the performance results measured and available so that I can compare implementations and make informed decisions.
As a user, I would like for the new Java Inference API to support RNNs so that the models I’ve trained can be deployed to our production environments.
As a user already familiar with MXNet, I’d like for the new API to be similar to existing implementations so that it’s easy for me to use.
Proposed Approach
The proposed implementation for the new Java API is to create a Java friendly wrapper around the existing Scala API. The Scala API is already fully implemented and is undergoing significant improvements (most notably simplifying the memory management of off-heap memory). By utilizing the existing Scala API, the development effort require for the new Java API is greatly decreased. Additionally, the Java API would automatically (or with minimal work) benefit from new features and code improvements allowing for development efforts to remain focused.
Since both Java and Scala are JVM languages, it is already possible for the Scala bindings to be called from Java code by loading the jar into the classpath. Due to differences in the languages, this process is currently very painful for users to implement. To improve upon this experience, a Java wrapper would be created which will call the Scala bindings. The wrapper would be designed so that it abstracts away the complexities of the Java/Scala interaction and is more idiomatic for the Java inferencing use case.
- Advantages
- Interaction with the native code is already done.
- The Scala API is already designed and decided. Implementing a wrapper limits design decisions which needs to be made and keeps the APIs consistent.
- Allows for development continue to be focused on a single JVM implementation which can be utilized by other JVM languages.
- The implementation, adding new features, maintenance would be greatly simplified.
- Implementation is not one way. In the future we maintain the ability to walk this decision back and go with another implementation.
- Disadvantages
- Interaction with the Scala code could be complicated due to differences in the languages. Known issues are forming Scala collections and leveraging default values. Changes to the Scala API or the introduction of a builder function seems like the most obvious solutions here.
- Some overhead in converting collections should be expected.
- The JAR files will be larger than they would be without Scala in the middle. Theoretically, this could be an issue for some memory constrained edge devices.
Known Difficulties
Converting Java collections into Scala collections - Scala and Java use different collections. Generally, these can be converted through the scala.collection.JavaConverters library. Ideally, this will be done automatically on behalf of the user. The Java methods should take Java collections, do the necessary conversion, then call the corresponding Scala method.
Java doesn’t support methods with default arguments - The current Scala implementation makes liberal use of default arguments. For class instantiation, a simple builder pattern will work. Class methods with default values will likely need to be overloaded.
Limited by existing Scala Inference API - The current Scala Inference API is lacking support for some models such as RNNs. Since this API will be utilized by the new Java Inference API, it will be necessary to improve and expand the Scala Inference API. This work can be done in parallel and should undergo it’s own design process.
Performance
Performance should be very similar to Scala. Since both are JVM languages doing inference will be calling the same byte code from Java as it is in Scala. The only known issue which will cause a performance difference is converting the Java collections into Scala collections. Preliminary testing with simple models shows negligible to nonexistent impact to performance. Java performance should be measured via Benchmark AI in a manner similar to how it's measured in Scala. More details on Scala benchmarks are available here.
Preliminary comparisons are a WIP.
Distribution
The new Java inference API can be distributed alongside the existing Scala API. Currently, the Scala API is distributed via a jar file using a Maven repository. There is ongoing work to automate this process and ideally this work will include the new Java API as well. The design for the Automated Scala Release is available here.
Class Diagram
Sequence Diagram
Java Inference API Design
The Java Inference API will be a wrapper around the high level Scala Inference interface. Here is an example of what the Java wrapper will look like for the Scala inference Predictor class.
/** * Implementation of prediction routines. * * @param modelPathPrefix Path prefix from where to load the model artifacts. * These include the symbol, parameters, and synset.txt * Example: file://model-dir/resnet-152 (containing * resnet-152-symbol.json, resnet-152-0000.params, and synset.txt). * @param inputDescriptors Descriptors defining the input node names, shape, * layout and type parameters * <p>Note: If the input Descriptors is missing batchSize * ('N' in layout), a batchSize of 1 is assumed for the model. * @param contexts Device contexts on which you want to run inference; defaults to CPU * @param epoch Model epoch to load; defaults to 0 */ Predictor(String modelPathPrefix, Array<DataDesc> inputDescriptors, Array[Context] Contexts, int epoch)
/** * Takes input as IndexedSeq one dimensional arrays and creates the NDArray needed for inference * The array will be reshaped based on the input descriptors. * * @param input: An Array of a one-dimensional array. An Array is needed when the model has more than one input. * @return Indexed sequence array of outputs */ Array <Array <Float>> predict(Array <Array <Float>> input)
/** * Predict using NDArray as input * This method is useful when the input is a batch of data * Note: User is responsible for managing allocation/deallocation of input/output NDArrays. * * @param inputBatch Array of NDArrays * @return Output of predictions as NDArrays */ Array <NDArray> predictWithNDArray(Array <NDArray> inputBatch)
Example Uses
List<float> doSinglePrediction(){ MXNetInferenceModel model = new MXNetInferenceModel(); model.load("/local/path/to/model"); List<float> singlePrediction = Arrays.asList(3.14159, .5772, 2.71828); model.predict(singlePrediction); return model.getPredictions().get(0); } List<List<float>> doBatchPrediction(){ MXNetInferenceModel model = new MXNetInferenceModel(); model.load("/local/path/to/model"); List<List<float>> batchPrediction = Arrays.asList( Arrays.asList(3.14159, .5772, 2.71828), Arrays.asList(.5772, 2.71828, 3.14159), Arrays.asList(2.71828, 3.14159, .5772)); model.predict(batchPrediction); return model.getPredictions(); }
Open Questions
How to deal with Option[T] field in Java when calling from Scala?
On Java side:
- Option 1: Create a wrapper class in Java that allow users to use Scala Option field smoothly, something like this.
On Scala side:
- Option 1: Use null to replace some field defined in Scala to match Java's need
- Option 2: Build override method in Scala to allow Java user to use them
- Option 3: Create builder on Scala side that allow Java user by pass the field for optional.
SCALA/JAVA INTEGRATION TIP
Construct interfaces in Java that define all types that will be passed between Java and Scala. Place these interfaces into a project that can be shared between the Java portions of code and the Scala portions of code. By limiting the features used in the integration points, there won’t be any feature mismatch issues. (referred from "scala-in-depth" page 242)
Possible Alternative Approaches
Writing a Java Inference API that directly calls the native code - Doing this would be designing and implementing a Java Inference API that will interact with the native code using the existing JNI code. The API would be designed to make Java inferencing simple and idiomatic. The existing JNI code could be shared by both the existing Scala API and the new Java Inference API. The biggest drawback to this approach is that it involves a significant amount of duplicate work that would be very difficult to maintain with current resources.
- Advantages
- No overhead from converting collections.
- No surprises from interacting with the Scala code.
- Would likely be the faster implementations since benchmarks generally show that Java outperforms Scala.
2. Disadvantages
- Duplication of efforts between this and the Scala API (this means reimplementing executor, ndarray, module, etc which is a significant effort).
- Will have to reimplement off-heap memory management.
- Added design effort to decide the Java API.
Adopt Java as the primary JVM language - This approach is basically to spend a very significant effort rewriting the entire Scala API into Java. After that was done we could begin adding support for other JVM languages using Java as a base and eventually the current Scala API could be deprecated. Obviously this involves a very significant upfront effort. Long-term it would be reasonable to expect improved performance across all JVM languages (since benchmarks generally show Java to outperform Scala) and it would likely be easier to add support for other JVM languages. The performance gains would likely be offset tremendously by the fact that most of the workload is done in the C++ native code and not in the JVM.
1. Advantages
Likely to see better performance across all JVM languages.
Easier to add support for more JVM languages in the future.
2. Disadvantages
Tremendous amount of upfront effort.
Scala API already exists and has been well received.
- Scala is a popular data science language
- Apache Spark is a Scala first language and is a popular analytics engine