Mark III Systems Blog

TensorTidbits: Tips for Navigating TensorFlow (GPU Management)

TensorFlow 2.4 is the latest release and contains a wide variety of changes and quality of life improvements for data scientists and developers. In our TensorTidbits series of blogs, we’ll go over some of the practical ways you can get more out of TensorFlow when you’re building and operationalizing your machine learning and deep learning models.

GPU Management

Previously in earlier versions of TensorFlow, users had to specify which version of TensorFlow you wanted to install, as there was a CPU-only and a GPU-only version and depending on whether GPUs were in scope or not for accelerated training and inference, there had to be an explicit choice made.  Since these versions had some very significant differences, it was relatively complicated, time consuming, and with some risk for users to maintain a codebase for both if you were developing and deploying on some machines with only CPUs and some with GPUs. Thankfully, with this updated version of TensorFlow, both CPU and GPU capabilities are included in the library and this previous issue has been abstracted away.

Finding Your GPU

This also means that users are not notified until runtime if there was an error finding a GPU in a given machine, so we’ve found the following command helpful in many of the things we’ve worked on. In order to help with the verification of the proper GPU setup, TensorFlow has included the following function:


This will return to you a boolean value depending on whether TensorFlow has access to your GPU or not.  It is helpful to test this BEFORE you begin your heavy data preparation so that you don’t have to wait until your 1 million images are ready to be used as training data before you realize you can’t load your graph because TensorFlow can’t see the GPU!

This function also allows the use of the “cuda_only” and “min_cuda_compute_capability” arguments which will allow you to check for CUDA capable only if you set `cuda_only=TRUE` and you can specify a major and minor version to indicate the minimum CUDA capability you are looking for.  Here is an example for searching for a CUDA-enabled GPU with at least a CUDA 7.0 capability.


    cuda_only=True, min_cuda_compute_capability=(7,0)

) For Future Releases

The above mentioned function will be deprecated in a later release so keep that in mind. Another way to get a list of the visible devices for TensorFlow is


which will return an array with all the GPU devices that TensorFlow can make use of.  Keep in mind that you need CUDA 11 and cuDNN 8 for TensorFlow 2.4 to work properly with your GPU.

You can also use this function to list the available CPU devices: