li-js / setGPU

Small Python library to automatically set CUDA_VISIBLE_DEVICES to the least loaded device on multi-GPU systems.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

setGPU

A small Python library that automatically sets CUDA_VISIBLE_DEVICES to the most vacant GPU (in terms of memory usage) on multi-GPU systems.

  • Installation: pip install git+https://github.com/li-js/setGPU.git@master
  • Usage: import setGPU before any import that will use a GPU like torch or tensorflow.

Dependencies

  • Jongwook Choi's gpustat library. It can be installed by pip install gpustat

Licensing

This code is in the public domain.

About

Small Python library to automatically set CUDA_VISIBLE_DEVICES to the least loaded device on multi-GPU systems.


Languages

Language:Python 100.0%