VeloxSeg/train.sh at master · JinPLu/VeloxSeg · 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
#!/usr/bin/env sh
set -eu

DATASET_NAME="${DATASET_NAME:-AutoPETII}"
MODEL_NAME="${MODEL_NAME:-VeloxSeg}"
TRAIN_CONFIG="${TRAIN_CONFIG:-./config/train_config_bs4.json}"
NUM_WORKERS="${NUM_WORKERS:-4}"
GPU_ID="${GPU_ID:-0}"

case "$DATASET_NAME" in
    AutoPETII)
        DEFAULT_MODEL_CONFIG="./config/models_config_autopetii.json"
        ;;
    Hecktor2022)
        DEFAULT_MODEL_CONFIG="./config/models_config_hecktor2022.json"
        ;;
    BraTS2021)
        DEFAULT_MODEL_CONFIG="./config/models_config_brats2021.json"
        ;;
    *)
        echo "Unsupported DATASET_NAME: $DATASET_NAME" >&2
        exit 2
        ;;
esac

MODEL_CONFIG="${MODEL_CONFIG:-$DEFAULT_MODEL_CONFIG}"

python -u ./run_train.py \
    --dataset_name "$DATASET_NAME" \
    --model_name "$MODEL_NAME" \
    --train_config "$TRAIN_CONFIG" \
    --model_config "$MODEL_CONFIG" \
    --num_workers "$NUM_WORKERS" \
    --gpu_id "$GPU_ID"