#!/bin/bash
# FPVE pruning experiment script for CelebA dataset

# Main training script
TRAIN_SCRIPT="python main_FPVE.py"

# Common parameters shared across all runs
COMMON_PARAMS="
               --load_dir your_model_path \
               --arch resnet18_half \
               --dataset celeba \
               --dataset_dir your_dataset_path \
               --target-attr Attractive \
               --batch-size 256 \
               --train_data_ratio 0.1 \
               --p-lr 0.1 \
               --ft_epochs 20 \
               --pop_init_rate 0.95 \
               --mutation_rate 0.05 \
               --pop_size 5 \
               --evolution_epoch 10 \
               --FPVE_fitness_data_ratio 1 \
               --fitness_mode ACC_FAIRNESS \
               --noadv_add_schedular \
               "

# if you want to use adversarial training, you can add the following parameters
# --adv_mode \
# --p-lr 0.0001 \
# --a-lr 0.0001 \
# --w 0.5 \


# Configuration parameters
PRUNING_RATIO=0.9
ITERATIVE_STEPS=9
SCALAR=0.1

# Seed and GPU mapping (parallel execution)
SEEDS=(2021 2022 2023 2024)  # e.g. (2021 2022 2023 2024)
GPUS=(0 1 2 3)     # e.g. (0 1 2 3)

# Function to start training
start_training() {
    local SEED=$1
    local GPU=$2
    local SCALAR=$3
    echo "Starting training with seed $SEED on GPU $GPU"

    SAVE_PATH=YOUR_DYNAMIC_SAVE_PATH

    $TRAIN_SCRIPT $COMMON_PARAMS \
        --pruning_ratio $PRUNING_RATIO \
        --iterative_steps $ITERATIVE_STEPS \
        --random_seed $SEED \
        --gpu $GPU \
        --scalar $SCALAR \
        --save_dir $SAVE_PATH &
}

# Main execution loop
for i in "${!SEEDS[@]}"; do
    SEED=${SEEDS[$i]}
    GPU=${GPUS[$i]}
    start_training "$SEED" "$GPU" "$SCALAR"
done

# Wait for all background processes to complete
wait
echo "All FPVE pruning tasks completed successfully."