python3获取nvidia GPU信息程序

最新推荐文章于 2023-11-17 16:56:05 发布

粼粼淇

最新推荐文章于 2023-11-17 16:56:05 发布

阅读量1.1k

点赞数 1

分类专栏： Linux Python

本文链接：https://blog.csdn.net/lingyunxianhe/article/details/118638629

版权

Python 同时被 2 个专栏收录

77 篇文章

订阅专栏

Linux

25 篇文章

订阅专栏

import subprocess
import numpy as np

def getcmdoutput(cmd):
    output=subprocess.getoutput(cmd)
    output =output.split('\n')
    return output


class GPUInfo:
    def getgpunum(self):
        cmd="nvidia-smi -L | grep GPU |wc -l"
        output=getcmdoutput(cmd)
        output=int(output[0])
        return output

    def gettotalgpumem(self):
        cmd = "nvidia-smi -q -d Memory | grep -A4 GPU | grep Total | grep -o '[0-9]\+'"
        output = getcmdoutput(cmd)
        output =[int(x) for x in output]
        return output

    def getavailgpumem(self):
        cmd = "nvidia-smi -q -d Memory | grep -A4 GPU | grep Free | grep -o '[0-9]\+'"
        output = getcmdoutput(cmd)
        output = [int(x) for x in output]
        return output

    def getavailgpumemratio(self,gpuid=None):
        totalgpumem=self.gettotalgpumem()
        availgpumem=self.getavailgpumem()
        availgpumemratio=np.array(availgpumem)/np.array(totalgpumem)
        if gpuid is not None:
            availgpumemratio=availgpumemratio[gpuid]
        return availgpumemratio

    def getbestgpuid(self):
        availgpumemratio=self.getavailgpumemratio()
        bestgpuid=np.argmax(availgpumemratio)
        return bestgpuid