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 [source]#
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()[source]#
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()[source]#
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] [source]#
Return a list of GPUs including unavailable devices.
- sleap.nn.system.get_available_gpus() List[PhysicalDevice] [source]#
Return a list of available GPUs.
- sleap.nn.system.get_current_gpu() PhysicalDevice [source]#
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] [source]#
Get the available memory on each GPU.
- Returns:
A list of the available memory on each GPU in MiB.
- sleap.nn.system.initialize_devices()[source]#
Initialize available physical devices as logical devices.
If preallocation was enabled on the GPUs, this will trigger memory allocation.
- sleap.nn.system.is_initialized(gpu: Optional[PhysicalDevice] = None) bool [source]#
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.use_cpu_only()[source]#
Hide GPUs from TensorFlow to ensure only the CPU is available.