python 版本的指定 GPUSwitcher

python 版本的指定 GPUSwitcher

Python 可以指定显卡使用 Cuda， 进行加速或者深度学习。通常通过 设置环境变量的方式。
环境变量为：“CUDA_VISIBLE_DEVICES”

import os


class GpuSwitcher(object):

    cuda_visible = "CUDA_VISIBLE_DEVICES"

    def __init__(self, cur_gpu_id):
        """ 设置当前的 GPU ID """
        if self.cuda_visible in os.environ.keys():
            self.ori_gpu_ids = os.environ[self.cuda_visible]
            os.environ["CUDA_VISIBLE_DEVICES"] = cur_gpu_id
            print(r'switch gpu: set current gpu id {}, origin gpu ids {}'.format(cur_gpu_id, self.ori_gpu_ids))
        else:
            print(r'switch gpu: CUDA_VISIBLE_DEVICES 不存在')

    def __del__(self):
        """ 恢复到原来的 GPU IDs """
        if self.cuda_visible in os.environ.keys():
            os.environ[self.cuda_visible] = self.ori_gpu_ids
            print(r'switch gpu: recover ori gpu ids {}'.format(self.ori_gpu_ids))
        else:
            print(r'switch gpu: CUDA_VISIBLE_DEVICES 不存在')


if __name__ == "__main__":
    GpuSwitcher(str(1))

GpuSwitcher 的作用是：在作用域中， 使用的就 cur_gpu_id。
当离开作用域时，就会恢复到原来的显卡状态。