脚本辅助消融实验

Python,Share,Linux 2024-08-17 34 次浏览 次点赞

常规来说,模型的train.py文件里会用argparse这个包,parser.add_argument()给模型规定一些数据集名称,还有一些超参的值。调试的时候一般在pycharm里手动改,平常训练,我一般就

python train.py

再或者后台挂起

nohup python train.py > outXXX.log &

但即使是后台挂起,不用维持SSH连接(一般是xshell SSH连集群),让你电脑可以关机(不会有人用自己电脑的显卡跑AI吧,这有点太伤了,比挖矿还难受。),但是在消融实验和对比实验的时候仍然不方便,这时候,用个脚本,让它自动检测显卡有没有被占用(每隔一段时间,我设的1分钟),然后没被占用就跑实验,占用就等待,这样可以不用随时盯着,即使是知道了大概时间,有时候也会忘掉,而且有时候它是半夜才停,别问我怎么知道的,我之前还定了闹钟,半夜爬起来改改参数,继续跑~

这时候,parser.add_argument()这个老管用了。

nohup python train.py --name=XXX --dataset=xxx --alpha_v=XXX > xxx.log &

这样就可以指定参数了,比如说上边--alpha_v=XXX明显就是$\alpha$这个超参的值设为多少。

现在开始看nvidia-smi

一般就双开左边的运行python train.py,右边运行watch nvidia-smi查看有啥变化,在python train.py前后。这一般就是判断显卡有无占用的判断依据。不同显卡可能稍微不一样。

比如说我用的集群

有占用时,最下边不显示

Every 2.0s: nvidia-smi                                                                                                                     20d938b479c0: Sat Aug 17 23:08:43 2024

Sat Aug 17 23:08:43 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| 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 P100-SXM2-16GB           Off |   00000000:B2:00.0 Off |                    0 |
| N/A   73C    P0            286W /  300W |    9048MiB /  16384MiB |     99%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage     
|=========================================================================================|
+-----------------------------------------------------------------------------------------+

没有占用时

Every 2.0s: nvidia-smi                                                                                                                     20d938b479c0: Sat Aug 17 23:08:43 2024
Sat Aug 17 23:08:43 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| 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 P100-SXM2-16GB           Off |   00000000:B2:00.0 Off |                    0 |
| N/A   73C    P0            286W /  300W |    9048MiB /  16384MiB |     99%      Default |
|                                         |                        |                  N/A 
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage     
|========================================================================================|No running processes found
+-----------------------------------------------------------------------------------------+

所以No running processes found就成了我的判断依据。

所以我写的run_train.sh脚本如下:

#!/bin/bash

task_list=(
    "nohup python train.py --name=XXX --dataset=XXX --alpha_v=xxx > XXX.log &"
    "nohup python train.py --name=XXX --dataset=XXX --alpha_v=xxx > XXX.log &"
    ……
)

check_gpu_usage() {
    local task_len=${#task_list[@]}
    local task_index=0

    while true; do
        nvidia_smi_output=$(nvidia-smi)
        
        if [[ $nvidia_smi_output == *"No running processes found"* ]]; then
            if (( task_index >= task_len )); then
                echo "All tasks completed successfully."
                exit 0
            else
                echo "No GPU processes found. Starting task $(( task_index + 1 )): ${task_list[task_index]}"
                eval "${task_list[task_index]}"
                (( task_index++ ))
                sleep 60
            fi
        else
            echo "GPU is currently in use. Waiting..."
            sleep 60
        fi
    done
}

check_gpu_usage

应该挺好看懂的,就是每隔一段时间检查gup使用情况,被占用,则等待,没有占用则按照顺序运行下一条命令,直至所有命令都运行完毕。

先创建run_train.sh,再编辑该文件,再赋予该文件执行的权限

[root@fmujieserver ~]# touch run_train.sh
[root@fmujieserver ~]# ls
anaconda-ks.cfg  frp  igbinary.sh  install.sh  original-ks.cfg  package.xml  run_train.sh  rustdesk  test4
[root@fmujieserver ~]# vim run_train.sh 
[root@fmujieserver ~]# chmod -x run_train.sh

最后

[root@fmujieserver ~]# nohup ./run_train.sh > out_sh.log &

这样,基本不用自己盯着了,毕竟一个实验跑挺久的。然后偶尔去服务器查看out_sh.log就行。


本文由 fmujie 创作,采用 知识共享署名 3.0,可自由转载、引用,但需署名作者且注明文章出处。

还不快抢沙发

添加新评论

召唤看板娘