经常要做模型开发训练的人一定对于GPU的查询不会陌生,实例如下:
详情数据如下:
Tue Aug 13 16:42:31 2024
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.80.02 Driver Version: 450.80.02 CUDA Version: 11.0 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla T4 On | 00000000:00:08.0 Off | 0 |
| N/A 69C P0 69W / 70W | 9389MiB / 15109MiB | 54% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| 0 N/A N/A 26725 C python3.8 6039MiB |
| 0 N/A N/A 27327 C python3 3347MiB |
+-----------------------------------------------------------------------------+
以往,启动新的模型训练任务或者是关闭老的模型训练任务大都是人工来直接执行nvidia-smi查询命令,之后根据GPU的实际占用情况来判断是否要执行对应操作的,这种方式非常常见,但是当实际项目开发中有很多任务等着去开发训练的时候,如何能够更加高效地利用服务器资源就是一个关键了,举个简单的例子,比如:我们的A模型正常结束训练的时间是凌晨3点半,这个时候我们基本都在睡梦中了,基本上不会说专门再起来去操作服务器启动新的模型训练了,到了第二天八点半工作的时候,再启动新的训练任务,粗略估计已经浪费了5个小时的时间。
这时候就在想,如果能够定时对GPU的占用情况进行查询,如果发现空闲资源的话那就自动启动新的训练任务,这样岂不是最佳的资源利用了吗?
这个想法的关键在于要对nvidia-smi查询命令的输出实现自动解析处理,提取对应的数据。
我们可以简单的采用文件存储的形式来保存每一次查询的结果,如下:
os.system("nvidia-smi >> gpu.txt")
之后提取对应的显存占用、利用率等情况:
info_list = [one.strip() for one in one_line.split("|") if one.strip()]
used_memory, total_memory = [i.strip()[:-3] for i in info_list[-2].split("/")]
utilization = info_list[-1].replace("Default","").strip()[:-1]
print(f"Total Memory: {total_memory}MiB")
print(f"Used Memory: {used_memory}MiB")
print(f"Utilization: {utilization}%")
实例输出如下:
接下来提取当前系统环境下的训练进程信息,如下:
process_matches = process_pattern.findall(gpuInfo)
process_info = [[int(pid), name, f"{mem}MiB"] for pid, name, mem in process_matches]
for info in process_info:
print(info)
实例输出如下所示:
接下来就可以设定定时操作了,如下:
res_list=[]
while True:
now_ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
used_memory, utilization, process_info = getGpuInfo()
one_time_list = [now_ts, used_memory, utilization, process_info]
res_list.append(one_time_list)
print("res_list_length: ", len(res_list))
with open("res_list.json","w") as f:
f.write(json.dumps(res_list))
time.sleep(2)
if used_memory <= system_threshold:
......
我这里的逻辑是我自己项目开发中的需求,未必是你们正好需要的,我这里的逻辑就是定时执行查询操作,判断总进程占用的显存量是否小于指定阈值,如果是的话则会自动启动新的训练任务,你可以预先把自己的训练任务放在任务队列中等待提取即可。
在我上面的程序中,按照时间先后对于查询得到的各项指标数据进行了存储,这个数据方便后续用来做可视化,比如:我们想要了解模型训练期间GPU的利用率变化情况,如下:
with open("res_list.json") as f:
res_list=json.load(f)
res_list=sorted(res_list,key=lambda e:e[0])
plt.clf()
plt.figure(figsize=(16,6))
plt.subplot(121)
used_memory=[float(one[1]) for one in res_list]
plt.plot(used_memory)
plt.title("used_memory cruve")
plt.subplot(122)
utilization=[float(one[2]) for one in res_list]
plt.plot(utilization)
plt.title("utilization cruve")
plt.savefig("gpu.png")
可视化结果如下:
可以看到:整个训练过程中,GPU的利用率变化幅度还是比较大的。
感兴趣的话也都可以自行结合自己的实际需求来进行改造开发。