10. GPU Support in TensorFlow for NVIDIA and MacOs

In TensorFlow, you can set the device (CPU or GPU) similarly to how you do it in PyTorch. TensorFlow will automatically use an available GPU if it's present, but you can explicitly check and set the device. Here's how you can do it:

import tensorflow as tf

# Check if any GPUs are available
if tf.config.list_physical_devices('GPU'):
    device = '/GPU:0'
elif tf.config.list_physical_devices('MPS'):  # For Apple Silicon (M1/M2) GPUs
    device = '/MPS:0'
else:
    device = '/CPU:0'

print(f"Using device: {device}")

# Example usage with TensorFlow's `with` context manager to specify the device
with tf.device(device):
    # Your model or tensor operations go here
    a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
    b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
    c = tf.matmul(a, b)
    print(c)

Explanation

1. tf.config.list_physical_devices('GPU'):

   • This checks if there are any available GPUs on the system. If there are, it sets the device to '/GPU:0'.

2. tf.config.list_physical_devices('MPS'):

   • This is specific to macOS and checks for the availability of an MPS (Metal Performance Shaders) device, which is used by Apple Silicon GPUs (M1/M2 chips). If available, it sets the device to '/MPS:0'.

3. '/CPU:0':

   • If no GPU or MPS device is found, the code falls back to using the CPU.

4. tf.device(device):

   • This is used to explicitly set the device for TensorFlow operations within the with context.

Summary

• GPU: Used if a CUDA-compatible GPU is available.
• MPS: Used for Apple Silicon GPUs on macOS.
• CPU: Used if no GPU or MPS device is available.

Click here to read about PyTorch device selection

Related Problems:

- What is Tensor Beginner