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

[rctay]

This template is for miscellaneous issues not covered by the other issue categories.

Hi, is there a way to check via Tensorflow Java the number of / or any CUDA devices available? For example, in Python with Torch one could do

import torch
torch.cuda.is_available()

Thanks in advance!

[karllessard]

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);
    }
  }
}