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。
当离开作用域时,就会恢复到原来的显卡状态。