脚本辅助消融实验

常规来说，模型的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就行。