指令小工具

由于实验室GPU服务器经常爆满，所以写了个GPU监控脚本来监控显卡使用情况，当有空闲显卡时立马把程序挂载上去
持续更新，shell监控进程 以及利用pushplus推送消息。由于苹果邮箱最多设置15分钟抓取一次，或者打开app才会抓取刷新，原来的邮箱推送（相关代码可以看厦大自动健康打卡脚本）具有一定的延迟，所以写了一个微信推送结合watch的微信消息提醒，很完美的实时提醒。

显卡监控

首先获取GPU的状态，其中显存占用在第三栏所以直接遍历序号2获得显存占用情况，当显存占用小于总显存一半，则记录下来（同时记录剩余显存和显卡号，便于后续训练指令）

1
2
3
4
5
6
7
8
9
10
11
12
def gpu_info():
    gpu_status = os.popen('nvidia-smi|grep %').read().split('|')
    gpu_memory = 0 # 记录剩余显存
    gpu_ids = [] # 记录可用显卡编号
    for i in range(2, int(len(gpu_status)/1),4):
        curr,totall = gpu_status[i].split('/')
        curr = int(curr.split('M')[0].strip())
        totall = int(totall.split('M')[0].strip())
        if totall-curr > totall/2:
            gpu_memory += totall-curr
            gpu_ids.append(int(i/4))
    return gpu_memory, gpu_ids

然后在主程序中写个死循环，每半分钟监控一次GPU当有满足训练所需显存量时通过执行shell脚本将程序挂载上去
1
2
3
4
5
6
7
8
9
10
11
if __name__ == '__main__':
    gpu_memory, gpu_ids = gpu_info()
    while True:
        if gpu_memory > 12000:
            print(time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime()) + " start gpu_memory=%d" % gpu_memory)
            break
        print(time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime()) + " waiting gpu_memory=%d" % gpu_memory)
        time.sleep(30)
        gpu_memory, gpu_ids = gpu_info()
    gpu_str = ','.join(str(i) for i in gpu_ids)
    os.system('sh train.sh %s %d'%(gpu_str, len(gpu_ids)))

最后给出执行训练的shell脚本，这里直接传参用哪几块显卡训练
1
2
3
4
5
6
7
8
9
#!/bin/bash

cd /home/hjz/work/TransUNet-main

DATA_DIR = /home/hjz/data/xerophthalmia

echo "start train"

CUDA_VISIBLE_DEVICES=$1 nohup python train.py -dataset own --num_classes 2 n_gpu $2 --vit_name ViT-B_16 --root_path DATA_DIR

Shell监控进程

因为训练和测试都采用shell脚本执行，所以在shell脚本下面可加入如下几行代码在训练后脚本持续监控训练进程 如果出现进程被kill的情况就会发送信息。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 发送开始消息
python SetEmail.py --messege 'xe_score程序开始' # 邮箱推送
python SetMessegeToWechat.py --title '程序变动' --content '程序开始' # 微信推送
sleep 5s
while true
do
  count=`ps -ef |grep hjz|grep python|grep $MODEL_NAME|wc -l` # 进程过滤 $MODEL_NAME 替换成当前训练进程指令中特有的字段。
  echo $count
  if [ 0 -eq $count ];then
    echo 'pass is not exist'
    # 发送进程终端消息
    python SetEmail.py --messege '进程中断' # 邮箱推送
    python SetMessegeToWechat.py --title '程序变动' --content '进程中断' # 微信推送
    break
  fi;
  sleep 60s # 每分钟监控一次
done

pushplus 提醒服务

  1. 首先登入网站扫码获取token
  2. 推送代码如下 SetMessegeToWechat.py
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    import requests
    import argparse
    import time

    parser = argparse.ArgumentParser()

    parser.add_argument('--title', type = str,default ='程序变动')
    parser.add_argument('--content', type = str,default ='训练成功')
    def send_notice(title,content):
        token = "你的token"
        url = f"http://www.pushplus.plus/send?token={token}&title={title}&content={content}&template=html"
        response = requests.request("GET", url)


    if __name__ == "__main__":
    	args = parser.parse_args()
    	send_notice(args.title, time.strftime("%Y-%m-%d %H:%M:%S  ", time.localtime()) + args.content)