GREB/google.sh at main · ustcllm/GREB · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
#!/bin/bash

# 定义类别列表
categories=(
    "Alabama"
    "Alaska"
    "Arizona"
    "Arkansas"
    "California"
    "Colorado"
    "Connecticut"
    "Delaware"
    "District_of_Columbia"
    "Florida"
    # "Georgia"
    # "Hawaii"
    # "Idaho"
    # "Illinois"
    # "Indiana"
    # "Iowa"
    # "Kansas"
    # "Kentucky"
    # "Louisiana"
    # "Maine"
    # "Maryland"
    # "Massachusetts"
    # "Michigan"
    # "Minnesota"
    # "Mississippi"
    # "Missouri"
    # "Montana"
    # "Nebraska"
    # "Nevada"
    # "New_Hampshire"
    # "New_Jersey"
    # "New_Mexico"
    # "New_York"
    # "North_Carolina"
    # "North_Dakota"
    # "Ohio"
    # "Oklahoma"
    # "Oregon"
    # "Pennsylvania"
    # "Rhode_Island"
    # "South_Carolina"
    # "South_Dakota"
    # "Tennessee"
    # "Texas"
    # "Utah"
    # "Vermont"
    # "Virginia"
    # "Washington"
    # "West_Virginia"
    # "Wisconsin"
    # "Wyoming"
)

# 定义CUDA设备组
devices=("0,1" "2,5" "4,3" "6,7")

# 当前正在运行的任务数
current_jobs=0

# 启动一个任务的函数
start_task() {
    local category=$1
    local device=$2

    echo "Starting task for category ${category} on devices ${device}"
    CUDA_VISIBLE_DEVICES="${device}" python script/run_ctr.py "${category}" &
}

# 启动任务的主循环
for category in "${categories[@]}"; do
    while [ "${current_jobs}" -ge 8 ]; do
        # 等待任何一个任务完成
        wait -n
        current_jobs=$((current_jobs - 1))
    done

    # 选择下一个可用设备组
    device=${devices[$((current_jobs / 2 % 4))]}

    start_task "${category}" "${device}"
    current_jobs=$((current_jobs + 1))
done

# 等待所有任务完成
wait

echo "All tasks have been trained."