python项目自动选择可用gpu

最新推荐文章于 2024-06-09 15:38:45 发布

无脑敲代码，bug漫天飞

最新推荐文章于 2024-06-09 15:38:45 发布

阅读量353

点赞数 1

分类专栏：编程文章标签： python 深度学习机器学习

本文链接：https://blog.csdn.net/qq_40671063/article/details/129622175

版权

编程专栏收录该内容

96 篇文章 9 订阅

订阅专栏

该代码段实现了一个功能，用于自动检测系统中NVIDIAGPU的空闲内存，并选择内存最大的GPU设置为CUDA_VISIBLE_DEVICES环境变量。通过nvidia-smi命令获取GPU内存信息，然后使用numpy找到最大空闲内存的GPU索引。

摘要由CSDN通过智能技术生成

code:

import os
import numpy as np

def choose_gpu():

    os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
    memory_gpu=[int(x.split()[2]) for x in open('tmp','r').readlines()]

    max_free = np.argmax(memory_gpu)
    print("choose gpu %d free %d MiB"%(max_free, memory_gpu[max_free]))

    os.environ['CUDA_VISIBLE_DEVICES']=str(max_free)
    os.system('rm tmp')

if __name__ == "__main__":
    choose_gpu()

参考：自动选择可用gpu - 简书