setGPU
A small Python library that automatically sets CUDA_VISIBLE_DEVICES
to the least-loaded GPU on multi-GPU systems and can be used by:
- Putting
import setGPUbefore any import that will use a GPU like Torch, TensorFlow, or JAX. - Defining an alias such as
alias setGPU='eval $(python3 -m setGPU)'and calling that to set the GPU in the shell before running another program that uses the GPU.
Installation
pip install git+https://github.com/bamos/setGPU.git
Dependencies
- Jongwook Choi's gpustat library (
pip install gpustat)
Licensing
This code is in the public domain.