Returns whether TensorFlow can access a GPU. Warning: if a non-GPU version of the package is installed, the function would also return False. Use `tf.test.is_built_with_cuda` to validate if TensorFlow was build with CUDA support. Args: cuda_only: limit the search to CUDA GPUs. min_
(cuda_only=False, min_cuda_compute_capability=None)
| 1385 | |
| 1386 | @tf_export("test.is_gpu_available") |
| 1387 | def is_gpu_available(cuda_only=False, min_cuda_compute_capability=None): |
| 1388 | """Returns whether TensorFlow can access a GPU. |
| 1389 | |
| 1390 | Warning: if a non-GPU version of the package is installed, the function would |
| 1391 | also return False. Use `tf.test.is_built_with_cuda` to validate if TensorFlow |
| 1392 | was build with CUDA support. |
| 1393 | |
| 1394 | Args: |
| 1395 | cuda_only: limit the search to CUDA GPUs. |
| 1396 | min_cuda_compute_capability: a (major,minor) pair that indicates the minimum |
| 1397 | CUDA compute capability required, or None if no requirement. |
| 1398 | |
| 1399 | Note that the keyword arg name "cuda_only" is misleading (since routine will |
| 1400 | return true when a GPU device is available irrespective of whether TF was |
| 1401 | built with CUDA support or ROCm support. However no changes here because |
| 1402 | |
| 1403 | ++ Changing the name "cuda_only" to something more generic would break |
| 1404 | backward compatibility |
| 1405 | |
| 1406 | ++ Adding an equivalent "rocm_only" would require the implementation check |
| 1407 | the build type. This in turn would require doing the same for CUDA and thus |
| 1408 | potentially break backward compatibility |
| 1409 | |
| 1410 | ++ Adding a new "cuda_or_rocm_only" would not break backward compatibility, |
| 1411 | but would require most (if not all) callers to update the call to use |
| 1412 | "cuda_or_rocm_only" instead of "cuda_only" |
| 1413 | |
| 1414 | Returns: |
| 1415 | True if a GPU device of the requested kind is available. |
| 1416 | """ |
| 1417 | |
| 1418 | def compute_capability_from_device_desc(device_desc): |
| 1419 | # TODO(jingyue): The device description generator has to be in sync with |
| 1420 | # this file. Another option is to put compute capability in |
| 1421 | # DeviceAttributes, but I avoided that to keep DeviceAttributes |
| 1422 | # target-independent. Reconsider this option when we have more things like |
| 1423 | # this to keep in sync. |
| 1424 | # LINT.IfChange |
| 1425 | match = re.search(r"compute capability: (\d+)\.(\d+)", device_desc) |
| 1426 | # LINT.ThenChange(//tensorflow/core/\ |
| 1427 | # common_runtime/gpu/gpu_device.cc) |
| 1428 | if not match: |
| 1429 | return 0, 0 |
| 1430 | return int(match.group(1)), int(match.group(2)) |
| 1431 | |
| 1432 | try: |
| 1433 | for local_device in device_lib.list_local_devices(): |
| 1434 | if local_device.device_type == "GPU": |
| 1435 | if (min_cuda_compute_capability is None or |
| 1436 | compute_capability_from_device_desc( |
| 1437 | local_device.physical_device_desc) >= |
| 1438 | min_cuda_compute_capability): |
| 1439 | return True |
| 1440 | if local_device.device_type == "SYCL" and not cuda_only: |
| 1441 | return True |
| 1442 | return False |
| 1443 | except errors_impl.NotFoundError as e: |
| 1444 | if not all(x in str(e) for x in ["CUDA", "not find"]): |
no test coverage detected