zea.init_device

zea.init_device(device='auto:1', backend='auto', hide_devices=None, allow_preallocate=True, verbose=True)[source]

Automatically selects a GPU or CPU device.

Useful to call at the start of a script to set the device for tensorflow, jax or pytorch. The function will select a GPU based on available memory, or fall back to CPU if no GPU is available.

Parameters:

  • backend (Optional[str]) – String indicating which backend to use. Can be ‘torch’, ‘tensorflow’, ‘jax’, ‘numpy’, None or “auto”. - When “auto”, the function will select the backend based on the KERAS_BACKEND environment variable. - For numpy this function will return ‘cpu’.

  • device (Union[str, int, list]) – device(s) to select. Examples: ‘cuda:1’, ‘gpu:2’, ‘auto:-1’, ‘cpu’, 0, or [0,1,2,3]. For more details see: get_device.

  • hide_devices (Union[int, list]) – device(s) to hide from the system. Examples: 0, or [0,1,2,3]. Can be useful when some GPUs have too little tensor cores to be useful for training, or when some GPUs are reserved for other tasks. Defaults to None, in which case no GPUs are hidden and all are available for use.

  • allow_preallocate (bool) – allow preallocation of memory. Used for jax and tensorflow.

  • verbose (bool) – print device selection. Defaults to True.

Returns:

selected device(s).

Return type:

device (str/int/list)