zea.init_device

zea.init_device(device='auto:1', backend='auto', hide_devices=None, hide_others=True, 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.

Generally, it is recommended to use init_device before importing any other library, since it will hide other GPUs from the libraries. The returned device string for a single GPU is ‘gpu:0’, since the others are hidden. Alternatively, you can set hide_others=False and manage the device yourself. In that case the returned device will contain the original physical GPU ids (e.g. 'gpu:2' for physical GPU 2).

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 before selection. 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.

  • hide_others (bool) – if True (default), unselected GPUs are hidden via CUDA_VISIBLE_DEVICES and selected GPUs are renumbered 0..N-1 in the returned device strings. If False, all GPUs remain visible and the returned device strings use the original physical GPU ids (e.g. 'gpu:2' for physical GPU 2).

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

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

Returns:

selected device string (e.g. 'gpu:0') or

a list of device strings (e.g. ['gpu:0', 'gpu:1']) when multiple GPUs were selected. Returns 'cpu' when no GPU is available or device='cpu' was requested.

Return type:

device (str | list[str])