Advice needed: how to check number of/any CUDA devices available? #614

Activity

1. karllessard commented on Jun 29, 2025

   

   karllessard

   on Jun 29, 2025

   Last edited by karllessard

   Collaborator

   More actions

   TensorFlow Java calls the C/C++ API of TensorFlow underneath. Many functions are wrapped by an idiomatic Java interface for user-friendliness, but that requires some manual work, hence we don't have a wrapper for all of them. Turns out that listing the devices falls in the category of not being wrapped (unless I missed it?), so you need to invoke directly the less user-friendly C interface via our JNI adapter.

   Something like that works out for me, it should be fairly simple to wrap this up as a utility method of the EagerSession object though, let me know if you are interested in doing it, otherwise I might give it a try:

   import org.tensorflow.internal.c_api.*;
   import static org.tensorflow.internal.c_api.global.tensorflow.*;
   
   public class Devices {
   
     public static void printDevices() {
       try (TF_Status status = TF_Status.newStatus()) {
         TFE_ContextOptions options = TFE_ContextOptions.newContextOptions();
         TFE_Context context = TFE_Context.newContext(options, status);
         status.throwExceptionIfNotOK();
   
         TF_DeviceList deviceList = TFE_ContextListDevices(context, status);
         status.throwExceptionIfNotOK();
   
         int numDevices = TF_DeviceListCount(deviceList);
         for (int i = 0; i < numDevices; i++) {
           String deviceName = TF_DeviceListName(deviceList, i, status).getString();
           status.throwExceptionIfNotOK();
           String deviceType = TF_DeviceListType(deviceList, i, status).getString();
           status.throwExceptionIfNotOK();
           System.out.println("Device " + i + ": " + deviceName + ", type: " + deviceType);
         }
         // TODO: make sure to release resources even when an exception has been thrown
         TF_DeleteDeviceList(deviceList);
         TFE_DeleteContext(context);
         TFE_DeleteContextOptions(options);
       }
     }
   }

   👍React with 👍1Reacted by winfriedgerlach