Using GPUs in TensorFlow for Java.

The TensorFlow Java API is distributed in the org.tensorflow:tensorflow maven package. As of TensorFlow 1.4.0, this package did not include GPU support by default. However, by making the GPU-enabled native libraries available to the JVM, GPUs can be utilized. This example demonstrates that.

Background

1. For the purposes of this demo, the "model" computes the square of a 2x2 float matrix. This model is serialized in graph.pb (generated by create_graph.py). That said, the description below applies equally to any TensorFlow model.

2. A simple Java program loads the graph and executes it. It also prints the device on which the operations execute to the console. To compile and execute the graph on CPU, run:

   mvn compile exec:java -Dexec.args="graph.pb"
   

   Doing so will result in something like this on the console:

   output: (MatMul): /job:localhost/replica:0/task:0/cpu:0
   input: (Placeholder): /job:localhost/replica:0/task:0/cpu:0
   TensorFlow version: 1.4.0
   
   Input : [[1.0, 2.0], [3.0, 4.0]]
   Output: [[7.0, 10.0], [15.0, 22.0]]
   

   Which shows that the MatMul operation to compute the output tensor was executed on cpu.

Using GPUs

To utilize a GPU for the matrix multiplication in the same program, the CUDA-enabled Java native libraries need to be made available to the JVM.

1. Download and extract the CUDA-enabled TensorFlow native library:

   curl -L "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-gpu-linux-x86_64-1.4.0.tar.gz" |
     tar -xz
   

   (For more detailed instructions see: https://www.tensorflow.org/install/install_java#using_tensorflow_with_jdk)

2. Make this native library available to the JVM. When using maven, this can be done using the MAVEN_OPTS environment variable:

   export MAVEN_OPTS="-Djava.library.path=${PWD}"
   mvn exec:java -Dexec.args="graph.pb"
   

This execution with the CUDA-enabled native libraries available to the JVM will result in something like this on the console:

output: (MatMul): /job:localhost/replica:0/task:0/gpu:0
input: (Placeholder): /job:localhost/replica:0/task:0/gpu:0
TensorFlow version: 1.4.0

Input : [[1.0, 2.0], [3.0, 4.0]]
Output: [[7.0, 10.0], [15.0, 22.0]]

In particular, note the /gpu:0 in this run, in contrast to the /cpu:0 in the previous one. This is an indication that the MatMul TensorFlow operation was executed on the GPU.

Troubleshooting

If the steps described above did not result in the GPU being used, some troubleshooting steps, try the following:

1. Ensure that your setup satisfies the NVIDIA requirements to run TensorFlow with GPU support

2. If you believe they do, then collect some debugging information to ensure that the TensorFlow Java library was able to enabled GPU support. Specifically:

   export MAVEN_OPTS="-Dorg.tensorflow.NativeLibrary.DEBUG=1 -Djava.library.path=${PWD}"
   mvn exec:java -Dexec.args="graph.pb"
   

   The full logs here may give you a hint. If you're at a loss, feel free to file an issue containing the full output of the command above and I could try and take a look.