#!/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 your_batch_size \
               --train_data_ratio your_train_data_ratio \
               --p-lr your_predictor_lr \
               --ft_epochs your_finetune_epochs \
               --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 \ # optional, use if you want to use a cos scheduler

# if you want to use adversarial training, you can add the following parameters
# --adv_mode \
# --a-lr your_adversary_lr \
# --w your_adversary_weight \
# --w_decay \ # optional, use if you want to use weight decay
# --use_projection \ # optional, use if you want to use projection

# Configuration parameters
PRUNING_RATIO=YOUR_PRUNING_RATIO  # Set your pruning ratio here
ITERATIVE_STEPS=YOUR_ITERATIVE_STEPS  # Set your iterative steps here
SCALAR=YOUR_SCALAR

# Seed and GPU mapping (parallel execution)
SEEDS=(YOUR_SEED_1 YOUR_SEED_2 YOUR_SEED_3)  # e.g. (2021 2022 2023 2024)
GPUS=(YOUR_GPU_1 YOUR_GPU_2 YOUR_GPU_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."