Tensorflow JNI development

From ElphelWiki
Revision as of 10:37, 5 March 2020 by Oleg (talk | contribs) (Install JAR to local Maven Repository)
Jump to: navigation, search


Why modify TF JNI?

  • Add TF features that are still missing in TF for Java like feeding directly from GPU memory thus saving time on back and forth CPU-GPU transfers if you, say, run data through a custom CUDA kernel first.


Notes on how to build TF JNI, where to modify JNI if needed, install to local maven and setup a project that will use modified native functions.

Based on Build TensorFlow 2.0 for Java on Windows article. Also this one - tensorflow/java/README.md.

These instructions are for Linux and old TensorFlow 1.15.0.

How to:

  • Build TF JNI - libtensorflow.jar, libtensorflow_jni.so and pom.xml
  • Add TF JAR to local Maven which will override the Central Maven Repository
  • Modify TF JNI functions
  • Create Elipse project

There's JavaCPP Presets project. Seems useless.


In Kubuntu:

Based on Feeding Tensorflow from GPU.


cd ~/git/tensorflow-1.15.0
./configure # do not forget CUDA
bazel build -c opt //tensorflow/java:tensorflow //tensorflow/java:libtensorflow_jni //tensorflow/java:pom

With TF, bazel tends to rebuild everything from scratch - takes a ton of time. Is it because it gets restarted after idle timeout or something else? A somewhat solution might be

At launch bazel starts its server which, to prevent it, add to ~/.bazelrc:
startup --max_idle_secs=0

Artifacts of interest are in bazel-bin/tensorflow/java/:

  • xml and jar will be taken care of by mvn command.
  • so will have to be in the library path. Link or copy to /usr/lib/ or go with "java -Djava.library.path=...". Example:
# /usr/lib is in the default java.library.path
sudo ln -sf ~/GIT/tensorflow-1.15.0/bazel-bin/tensorflow/java/libtensorflow_jni.so /usr/lib/

Install JAR to local Maven Repository

~/GIT/tensorflow-1.15.0$ mvn install:install-file -Dfile=bazel-bin/tensorflow/java/libtensorflow.jar -DpomFile=bazel-bin/tensorflow/java/pom.xml

How to uninstall maven local repo - and switch back to official versions from Maven Central - this link. Or remove unneeded stuff from ~/.m2/repository/org/tensorflow

After *_jni.so is linked and jar installed one can resume normal development. See below what to add to your project's pom.xml

Modify TF JNI functions

For example, one wants to create a new function in org.tensorflow.TensorFlow package. Then see inside:


Three places:

  • add native method to tensorflow/java/src/main/java/org/tensorflow/TensorFlow.java
  • add to header file tensorflow/java/src/main/native/tensorflow_jni.h
  • add to c file tensorflow/java/src/main/native/tensorflow_jni.cc

Rebuild and Reinstall.

The native header files seem to be regenerated but I haven't tested if they are actually used (need to test). In function naming - avoid underscores, e.g.:


Java Maven project in Eclipse

Nothing special.

  • Create a new maven project
  • Edit pom.xml:
  • Write code as usual

Basic example code


import org.tensorflow.TensorFlow;

public class tfhello{
	public static void main(String[] args){

A few words on TF in Maven Central repository


Record in pom.xml:


Archive contains Java classes.


Record in pom.xml:


Archive contains native library:

│   └── maven
│       └── org.tensorflow
│           └── libtensorflow_jni_gpu
│               ├── pom.properties
│               └── pom.xml
└── org
    └── tensorflow
        └── native
            ├── linux-x86_64
            │   ├── libtensorflow_framework.so.1
            │   ├── libtensorflow_jni.so
            │   ├── LICENSE
            │   └── THIRD_PARTY_TF_JNI_LICENSES
            └── windows-x86_64
                ├── LICENSE
                └── tensorflow_jni.dll