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/
Python跑pytorch程序抢占公共GPU自动运行脚本_第1张图片

当然,如果网络足够好,可以直接利用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会很难受,不信自行尝试
效果图:

你可能感兴趣的:(python,pytorch,开发语言)