os.environ[‘CUDA_VISIBLE_DEVICES‘]使用指定的GPU及GPU显存
程序员文章站
2024-03-24 10:21:16
...
if cuda_idx >= 0:
os.environ['CUDA_VISIBLE_DEVICES'] = str(cuda_idx) # cuda_idx在config中设置
服务器上有好几张显卡,别人可能在某几张卡上跑满了,这时需要你指定某张显卡来跑实验。使用上述功能,则代码只能对 str(cuda_idx)的显卡可见,并对所有可见的显卡依然按照数组从0开始编号。