138 lines
5.3 KiB
Bash
138 lines
5.3 KiB
Bash
#!/bin/bash
|
|
|
|
# =================================================================
|
|
# YOLO 模型批量并行训练脚本
|
|
# =================================================================
|
|
|
|
# --- 1. Conda 环境设置 ---
|
|
CONDA_BASE_PATH="/home/wkmgc/miniconda3" # <--- 在这里修改为您自己的 Conda 路径
|
|
CONDA_ENV_NAME="${SEG_CONDA_ENV:-seg_smp}" # 可用 SEG_CONDA_ENV=SMP bash yolo_train.sh 临时覆盖
|
|
|
|
# 初始化并激活 Conda 环境
|
|
if [ -f "${CONDA_BASE_PATH}/etc/profile.d/conda.sh" ]; then
|
|
source "${CONDA_BASE_PATH}/etc/profile.d/conda.sh"
|
|
conda activate "${CONDA_ENV_NAME}"
|
|
if [ $? -ne 0 ]; then
|
|
echo "错误: 激活 Conda 环境 '${CONDA_ENV_NAME}' 失败！"
|
|
exit 1
|
|
fi
|
|
echo "Conda 环境 '${CONDA_ENV_NAME}' 已成功激活。"
|
|
else
|
|
echo "错误: 找不到 conda.sh 脚本。请检查您的 CONDA_BASE_PATH 设置是否正确。"
|
|
exit 1
|
|
fi
|
|
|
|
echo "正在为当前终端会话设置 HTTP/HTTPS 代理..."
|
|
export https_proxy=http://127.0.0.1:1080
|
|
export http_proxy=http://127.0.0.1:1080
|
|
echo "代理已设置为 http://127.0.0.1:1080"
|
|
echo "如果模型权重不存在，将通过此代理进行下载。"
|
|
|
|
# --- 2. 模型与 GPU 配置 ---
|
|
# 根据您的 GPU 硬件和训练计划进行分组
|
|
# yolo_train.py 会自动使用分配给它的所有可见 GPU
|
|
GPUS_GROUP_0="0"
|
|
GPUS_GROUP_1="1"
|
|
GPUS_GROUP_2="2"
|
|
GPUS_GROUP_3="3"
|
|
# TODO
|
|
GPUS_GROUP_4="0"
|
|
GPUS_GROUP_5="1"
|
|
GPUS_GROUP_6="2"
|
|
GPUS_GROUP_7="3"
|
|
|
|
# GPUS_GROUP_4="4"
|
|
# GPUS_GROUP_5="5"
|
|
# GPUS_GROUP_6="6"
|
|
# GPUS_GROUP_7="7"
|
|
|
|
# 从 yolo_config.py 中选择要训练的模型，并分配到不同的组
|
|
# 建议将计算量/模型大小相似的模型放在同一组
|
|
GROUP_0_MODELS=("YOLO11l-seg" "YOLOv8m-seg")
|
|
GROUP_1_MODELS=("YOLO11x-seg")
|
|
GROUP_2_MODELS=("YOLO11n-seg" "YOLO11s-seg" "YOLO11m-seg")
|
|
GROUP_3_MODELS=("YOLOv9e-seg")
|
|
GROUP_4_MODELS=("YOLOv9c-seg")
|
|
GROUP_5_MODELS=("YOLOv8x-seg")
|
|
GROUP_6_MODELS=("YOLOv8l-seg" "YOLO12-seg") # YOLOv9e-seg # 模型太大
|
|
GROUP_7_MODELS=("YOLOv8n-seg" "YOLOv8s-seg" ) # 训练完毕
|
|
|
|
# 1. 从 config.py 中读取 OUTPUTS_DIR 的值
|
|
OUTPUTS_DIR=$(python -c "from yolo_config import OUTPUTS_DIR; print(OUTPUTS_DIR)")
|
|
# 检查是否成功获取了 OUTPUTS_DIR
|
|
if [ -z "$OUTPUTS_DIR" ]; then
|
|
echo "OUTPUTS_DIR: $OUTPUTS_DIR"
|
|
echo "Error 1: Could not read OUTPUTS_DIR from yolo_config.py. Exiting."
|
|
echo "Error 2: Or the directory specified by OUTPUTS_DIR does not exist. Please create it first."
|
|
# exit 1
|
|
fi
|
|
# 2. 定义带有时间戳的日志目录名
|
|
LOG_DIR_NAME="yolo_train_logs_parallel_$(date +%Y-%m-%d_%H-%M-%S)"
|
|
# 3. 拼接成最终的完整路径
|
|
LOG_DIR="$OUTPUTS_DIR/$LOG_DIR_NAME"
|
|
mkdir -p "${LOG_DIR}"
|
|
echo "所有模型的日志将保存在 ./${LOG_DIR}/ 目录中。"
|
|
echo "----------------------------------------------------"
|
|
|
|
|
|
# --- 3. 训练执行函数 ---
|
|
# 定义一个函数来启动一组训练，以避免代码重复
|
|
start_training_group() {
|
|
# 使用 nameref (引用) 来传递数组
|
|
local -n models=$1
|
|
local gpus=$2
|
|
local group_name=$3
|
|
|
|
echo ">>> 准备启动 ${group_name} 的训练任务 (后台运行)..."
|
|
# 遍历指定组中的所有模型
|
|
for model_key in "${models[@]}"; do
|
|
echo " -> 正在后台启动模型: ${model_key} on GPUs: ${gpus}"
|
|
# 使用 '&' 将命令放入后台运行
|
|
# 通过 --model 参数将模型名称传递给 yolo_train.py
|
|
CUDA_VISIBLE_DEVICES=${gpus} python yolo_train.py --model "${model_key}" > "${LOG_DIR}/${model_key}.log" 2>&1 &
|
|
echo " - 模型 ${model_key} 已在后台启动。日志文件: ${LOG_DIR}/${model_key}.log"
|
|
echo " - 等待 30 秒，确保 GPU 资源稳定分配..."
|
|
sleep 30
|
|
done
|
|
echo ">>> ${group_name} 的所有模型均已启动。"
|
|
echo "----------------------------------------------------"
|
|
}
|
|
|
|
# --- 4. 依次启动所有训练任务 ---
|
|
# 脚本将快速地按顺序启动每一组任务到后台
|
|
start_training_group GROUP_0_MODELS "${GPUS_GROUP_0}" "第零组" # TODO
|
|
start_training_group GROUP_1_MODELS "${GPUS_GROUP_1}" "第一组"
|
|
start_training_group GROUP_2_MODELS "${GPUS_GROUP_2}" "第二组"
|
|
start_training_group GROUP_3_MODELS "${GPUS_GROUP_3}" "第三组"
|
|
|
|
|
|
# --- 5. 等待所有后台任务完成 ---
|
|
echo ""
|
|
echo "--- 所有模型均已在后台启动。现在等待所有训练任务完成... ---"
|
|
# 'wait' 命令会暂停脚本，直到所有由此脚本启动的后台子进程全部执行完毕
|
|
wait
|
|
echo "--- 所有后台训练任务已全部完成！ ---"
|
|
|
|
# 适用于4卡版本
|
|
start_training_group GROUP_4_MODELS "${GPUS_GROUP_4}" "第四组" # TODO
|
|
start_training_group GROUP_5_MODELS "${GPUS_GROUP_5}" "第五组"
|
|
start_training_group GROUP_6_MODELS "${GPUS_GROUP_6}" "第六组"
|
|
start_training_group GROUP_7_MODELS "${GPUS_GROUP_7}" "第七组"
|
|
# --- 5. 等待所有后台任务完成 ---
|
|
echo ""
|
|
echo "--- 所有模型均已在后台启动。现在等待所有训练任务完成... ---"
|
|
# 'wait' 命令会暂停脚本，直到所有由此脚本启动的后台子进程全部执行完毕
|
|
wait
|
|
echo "--- 所有后台训练任务已全部完成！ ---"
|
|
|
|
# --- 6. 退出脚本 ---
|
|
# 训练完成，取消激活 Conda 环境
|
|
echo "训练流程结束。"
|
|
# --- 取消网络代理设置 ---
|
|
echo "正在取消当前终端会话的 HTTP/HTTPS 代理设置..."
|
|
unset https_proxy
|
|
unset http_proxy
|
|
echo "代理已取消。"
|
|
conda deactivate
|
|
echo "取消激活 Conda 环境。"
|