#!/bin/bash
# This is a demo script for FPVE pruning experiments

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

# Common parameters shared across all runs
COMMON_PARAMS="
               --dataset cub200 \
               --data_root your_data_root_path \
               --arch resnet50 \
               --ckpt_load_dir your_model_path \
               --opt_name sgd \
               --lr 0.002 \
               --lr_scheduler_name cosineannealinglr \
               --batch_size 8 \
               --ft_epochs 50 \
               --pop_init_rate 0.95 \
               --mutation_rate 0.05 \
               --pop_size 5 \
               --evolution_epoch 10 \
               --FPVE_fitness_data_ratio 1 \
               --fitness_mode ACC_FAIRNESS \
               --fairness_eval_flag \
               --fairness_type max_min \
               "

# Pruning configuration
PRUNING_RATIO=0.2  # Set your pruning ratio here
ITERATIVE_STEPS=1  # Set your iterative steps here
SCALAR=0.01  # Set your scalar value here

# Seed and GPU mapping (parallel execution)
SEEDS=(2021 2022 2023 2024)
GPUS=(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 \
        --ckpt_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."