sleap.nn.system

Utilities for working with the physical system (e.g., GPUs).

This module mostly provides convenience functions for changing the state of the runtime environment by wrapping tf.config module functions.

sleap.nn.system.best_logical_device_name() → str

Return the name of the best logical device for performance.

This is particularly useful to use with tf.device() for explicit tensor placement.

Returns:

The name of the first logical GPU device if available, or alternatively the CPU.

Notes

This will initialize the logical devices if they were not already!

sleap.nn.system.disable_preallocation()

Disable preallocation of full GPU memory on all available GPUs.

This enables memory growth policy so that TensorFlow will not pre-allocate all available GPU memory.

Preallocation can be more efficient, but can lead to CUDA startup errors when the memory is not available (e.g., shared, multi-session and some *nix systems).

See also: enable_gpu_preallocation

sleap.nn.system.enable_preallocation()

Enable preallocation of full GPU memory on all available GPUs.

This disables memory growth policy so that TensorFlow will pre-allocate all available GPU memory.

Preallocation can be more efficient, but can lead to CUDA startup errors when the memory is not available (e.g., shared, multi-session and some *nix systems).

See also: disable_gpu_preallocation

sleap.nn.system.get_all_gpus() → List[PhysicalDevice]

Return a list of GPUs including unavailable devices.

sleap.nn.system.get_available_gpus() → List[PhysicalDevice]

Return a list of available GPUs.

sleap.nn.system.get_current_gpu() → PhysicalDevice

Return the current (single) GPU device.

Returns:

The tf.config.PhysicalDevice for the available GPU.

If no GPUs are available, returns None.

Raises:

ValueError – If multiple GPUs are available.

sleap.nn.system.get_gpu_memory() → List[int]

Get the available memory on each GPU.

Returns:

A list of the available memory on each GPU in MiB.

sleap.nn.system.initialize_devices()

Initialize available physical devices as logical devices.

If preallocation was enabled on the GPUs, this will trigger memory allocation.

sleap.nn.system.is_gpu_system() → bool

Return True if the system has discoverable GPUs.

sleap.nn.system.is_initialized(gpu: PhysicalDevice | None = None) → bool

Check if a physical GPU has been initialized without triggering initialization.

Parameters:

gpu – The GPU to check for initialization. If None, defaults to the current GPU.

Returns:

True if the GPU is initialized.

Notes

Once initialized, the GPU cannot be hidden or change its memory policy.

Initialization happens when a tf.config.LogicalDevice is created on the physical device, typically when the first tensor op runs.

Checking if the GPU is initializing by querying the logical devices will trigger initialization, so this method provides an easy way of checking without modifying the system state.

sleap.nn.system.summary()

Print a summary of the state of the system.

sleap.nn.system.use_cpu_only()

Hide GPUs from TensorFlow to ensure only the CPU is available.

sleap.nn.system.use_first_gpu()

Make only the first GPU available to TensorFlow.

sleap.nn.system.use_gpu(device_ind: int)

Make a single GPU available to TensorFlow.

Parameters:

device_ind – Index of the GPU within the list of system GPUs.

sleap.nn.system.use_last_gpu()

Make only the last GPU available to TensorFlow.