search

tensorflow / java

Java bindings for TensorFlow
Apache License 2.0
832 stars 202 forks source link

Add instructions on how to support GPU with TFJava 1.0.0 #565

Closed karllessard closed 3 weeks ago

crutcher commented 3 weeks ago

This section in CONTRIBUTING.md still describes how to build tensorflow-core-platform-gpu; which seems to also somehow be defined in the poms?

To build for GPU, pass -Djavacpp.platform.extension=-gpu to maven. If you want to use TensorFlow Java with unsupported GPUs, set the environment variable TF_CUDA_COMPUTE_CAPABILITIES, or configure it in a bazel rc file (i.e. build --action_env TF_CUDA_COMPUTE_CAPABILITIES="6.1").
Craigacp commented 3 weeks ago

javacpp.platform.extension is not related to tensorflow-core-platform-gpu, the use of platform in both names doesn't imply anything. Setting the extension triggers the GPU build, but it makes a tensorflow-core-native:linux-x86_64-gpu artifact.

karllessard commented 3 weeks ago

To do a custom native build that supports GPU (with the GPU of your choice) you need to pass -Djavacpp.platform.extension=-gpu in addition to -Pnative-build

(I just noticed that the doc actually says -Dnative-build but -Pnative-build would also work and sounds better, I'll update that as well)