88 lines
2.2 KiB
Bash
Executable File
88 lines
2.2 KiB
Bash
Executable File
#!/bin/bash
|
||
# BEVFusion 多任务训练脚本（检测 + 分割）
|
||
|
||
set -e # 遇到错误立即退出
|
||
|
||
echo "======================================"
|
||
echo "BEVFusion 多任务训练脚本"
|
||
echo "任务：3D目标检测 + BEV地图分割"
|
||
echo "======================================"
|
||
echo ""
|
||
|
||
# 配置参数
|
||
NUM_GPUS=8
|
||
CONFIG="configs/nuscenes/multitask/fusion-det-seg-swint.yaml"
|
||
BACKBONE_PRETRAIN="pretrained/swint-nuimages-pretrained.pth"
|
||
LIDAR_PRETRAIN="pretrained/lidar-only-det.pth"
|
||
|
||
# 检查预训练模型
|
||
echo "检查预训练模型..."
|
||
if [ ! -f "$BACKBONE_PRETRAIN" ]; then
|
||
echo "错误：找不到 $BACKBONE_PRETRAIN"
|
||
echo "请先下载预训练模型"
|
||
exit 1
|
||
fi
|
||
|
||
if [ ! -f "$LIDAR_PRETRAIN" ]; then
|
||
echo "错误：找不到 $LIDAR_PRETRAIN"
|
||
echo "请先下载预训练模型"
|
||
exit 1
|
||
fi
|
||
echo "✓ 预训练模型检查通过"
|
||
echo ""
|
||
|
||
# 检查配置文件
|
||
echo "检查配置文件..."
|
||
if [ ! -f "$CONFIG" ]; then
|
||
echo "错误：找不到配置文件 $CONFIG"
|
||
exit 1
|
||
fi
|
||
echo "✓ 配置文件检查通过"
|
||
echo ""
|
||
|
||
# 检查GPU
|
||
echo "检查GPU..."
|
||
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
|
||
echo "可用GPU数量: $GPU_COUNT"
|
||
if [ $GPU_COUNT -lt $NUM_GPUS ]; then
|
||
echo "警告：可用GPU数量($GPU_COUNT)小于配置数量($NUM_GPUS)"
|
||
echo "将使用 $GPU_COUNT 个GPU"
|
||
NUM_GPUS=$GPU_COUNT
|
||
fi
|
||
echo ""
|
||
|
||
# 显示训练配置
|
||
echo "======================================"
|
||
echo "训练配置："
|
||
echo " 配置文件: $CONFIG"
|
||
echo " GPU数量: $NUM_GPUS"
|
||
echo " Camera Backbone预训练: $BACKBONE_PRETRAIN"
|
||
echo " LiDAR预训练: $LIDAR_PRETRAIN"
|
||
echo " 任务: 3D检测 + BEV分割"
|
||
echo "======================================"
|
||
echo ""
|
||
|
||
# 询问确认
|
||
read -p "是否开始训练？(y/n) " -n 1 -r
|
||
echo
|
||
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
|
||
echo "训练已取消"
|
||
exit 0
|
||
fi
|
||
|
||
# 开始训练
|
||
echo ""
|
||
echo "开始训练..."
|
||
echo ""
|
||
|
||
torchpack dist-run -np $NUM_GPUS python tools/train.py \
|
||
$CONFIG \
|
||
--model.encoders.camera.backbone.init_cfg.checkpoint $BACKBONE_PRETRAIN \
|
||
--load_from $LIDAR_PRETRAIN
|
||
|
||
echo ""
|
||
echo "======================================"
|
||
echo "训练完成！"
|
||
echo "======================================"
|
||
|