Find usable CUDA devices¶
If you want to run several experiments at the same time on your machine, for example for a hyperparameter sweep, then you can
use the following utility function to pick GPU indices that are “accessible”, without having to change your code every time.
from lightning.pytorch.accelerators import find_usable_cuda_devices
# Find two GPUs on the system that are not already occupied
trainer = Trainer(accelerator="cuda", devices=find_usable_cuda_devices(2))
from lightning.fabric.accelerators import find_usable_cuda_devices
# Works with Fabric too
fabric = Fabric(accelerator="cuda", devices=find_usable_cuda_devices(2))
This is especially useful when GPUs are configured to be in “exclusive compute mode”, such that only one process at a time is allowed access to the device.
This special mode is often enabled on server GPUs or systems shared among multiple users.
|