python3获取nvidia GPU信息程序

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