autocuda - Auto choose the cuda device having the largest free memory in Pytorch

Usage

Install

pip install autocuda

ready to use

from autocuda import auto_cuda_info, auto_cuda_index, auto_cuda, auto_cuda_name

cuda_info_dict = auto_cuda_info()

cuda_device_index = auto_cuda_index()  # return cuda index having largest free memory. return 'cpu' if not cuda
# os.environ['CUDA_VISIBLE_DEVICES'] = [str(cuda_device_index)]

cuda_device = auto_cuda()
# model.to(cuda_device) # assume you have inited your pytorch model

cuda_device_name = auto_cuda_name()
print('Choosing cuda device: {}'.format(cuda_device_name))

Copyright

This package is based on https://github.com/QuantumLiu/tf_gpu_manager with MIT license