Advice needed: how to check number of/any CUDA devices available?
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!
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);
}
}
}