Python跑pytorch程序抢占公共GPU自动运行脚本

问题描述

当我们有一个服务器，服务器上面有4-5个GPU，那么我们需要时刻看哪个GPU空着，当发现服务器空闲了，我们就可以跑自己的深度学习了。

然而，人盯着总是费时费力的，所以可以让Python看到哪个GPU空闲就插进去吗？
进行下面步骤即可。

第一步、安装GPU信息查看包

名字为：nvidia_ml_py-12.535.133-py3-none-any.whl，
下载地址：https://pypi.org/project/nvidia-ml-py/

当然，如果网络足够好，可以直接利用pip安装：

pip install nvidia-ml-py

第二步、编码select_gpu

即判定抢占GPU的代码

import torch
from pynvml import *
import time
import sys


# count传GPU个数，threshold是阈值，低于此阈值说明GPU是空闲的，second是每几秒进行继续轮训
def select_gpu(count=torch.cuda.device_count(), threshold=1024, second=5):
    nvmlInit()
    if count == 0:
        return 'cpu'
    # 需要在多个GPU轮训找出空闲的GPU
    current = 0
    while True:
        # 检查当前GPU是否可用
        handle = nvmlDeviceGetHandleByIndex(current)
        info = nvmlDeviceGetMemoryInfo(handle)
        used_memory = info.used // (1024 * 1024)
        if used_memory < threshold:  # 如果刺入小于阈值的内存，那么说明此GPU并没有被占用，可抢占
            sys.stderr.write(
                f'此时GPU{current}使用内存为[{used_memory}MB]，低于阈值[{threshold}]才可抢占----GPU{current}可抢占，将抢占GPU：{current}号GPU----{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())}\n')
            nvmlShutdown()
            return current
        else:
            sys.stderr.write(
                f'此时GPU{current}使用内存为[{used_memory}MB]，低于阈值[{threshold}]才可抢占----GPU{current}不可抢占，继续轮训----{time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())}\n')
        time.sleep(second)
        current = (current + 1) % count

select_gpu共三个参数
select_gpu 用于抢占GPU
参数1：count传GPU个数，默认传torch.cuda.device_count()，即存在的GPU个数
参数2：threshold是阈值，低于此阈值说明GPU是空闲的，默认1024MB [GPU一般什么都不跑，也会被占用几十MB]
参数3：second是每几秒进行继续轮训，默认5秒

返回值：为可以选用GPU的编号。

第三步、编写运行深度学习代码，以YOLOv8为例

from ultralytics import YOLO
from select_gpu_util import select_gpu
import sys

if __name__ == '__main__':
    sys.stderr.write('程序已运行\n')
    # select_gpu 用于抢占GPU
    # 参数1：count传GPU个数，默认传torch.cuda.device_count()，即存在的GPU个数
    # 参数2：threshold是阈值，低于此阈值说明GPU是空闲的，默认1024MB [GPU一般什么都不跑，也会被占用几十MB]
    # 参数3：second是每几秒进行继续轮训，默认5秒
    device = select_gpu(threshold=3000)  # 获取能够抢占的GPU
    model = YOLO("ultralytics/cfg/models/v8/yolov8.yaml")  # build a new model from scratch
    # Train the model
    if sys.platform.startswith('win'):  # Windows环境下
        results = model.train(data="YAML/xxx.yaml", epochs=200, device=device, workers=0, batch=1)
    else:  # Linux系统环境下
        results = model.train(data="YAML/xxx.yaml", epochs=200, device=device, workers=2, batch=8)

# 温馨提示：启动命令为 nohup python -u train_shwd+GAM.py >shwdlog/gamdaoer.log &
#                                 -u的作用是不设置缓冲区，让所有的文本直接输出到log

第四步、运行代码，查看运行效果

运行一般让程序不挂起、并且后台运行，运行命令为：

nohup python -u train_shwd+GAM.py >shwdlog/gamdaoer.log &

-u的作用是不设置缓冲区，让所有的文本直接输出到log，若不-u会很难受，不信自行尝试
效果图：