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

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

# Common parameters that would be shared across all runs
COMMON_PARAMS="
               --dataset your_dataset \
               --data_root your_data_root_path \
               --arch your_arch \
               --ckpt_load_dir your_model_path \
               --opt_name your_optimizer \
               --lr your_lr \
               --lr_scheduler_name your_lr_scheduler \
               --batch_size your_batch_size \
               --ft_epochs your_epochs \
               --fairness_eval_flag \
               --fairness_type max_min \
               --layer_wise_imp \
               ... \
               "
               
# Pruning methods to test (sequential execution)
PRUNE_METHODS=("method1" "method2" "method3") # e.g. ("l1" "FPGM" ...)

# Seed and GPU mapping (parallel execution within each method)
SEEDS=(YOUR_SEED_1 YOUR_SEED_2 YOUR_SEED_3)  # e.g. (2021 2022 2023)
GPUS=(YOUR_GPU_1 YOUR_GPU_2 YOUR_GPU_3)      # e.g. (0 1 2)

# Pruning configuration
RATIO=YOUR_PRUNING_RATIO  # Set your pruning ratio here
ITERATIVE_STEPS=YOUR_ITERATIVE_STEPS  # Set your iterative steps here

# Outer loop: Sequentially execute each pruning method
for METHOD in "${PRUNE_METHODS[@]}"; do
    echo "================================================"
    echo "Starting pruning method: $METHOD"
    echo "================================================"
    
    # Inner loop: Parallel execution for each seed
    for i in "${!SEEDS[@]}"; do
        SEED=${SEEDS[$i]}
        GPU=${GPUS[$i]}
        
        # Generate save directory for this run
        SAVE_DIR_PARAM="--ckpt_save_dir YOUR_SAVE_PATH"
        
        echo "Starting $METHOD with seed $SEED on GPU $GPU"
        $TRAIN_SCRIPT $COMMON_PARAMS \
            --gpu $GPU \
            --random_seed $SEED \
            --pruning_ratio $RATIO \
            --iterative_steps $ITERATIVE_STEPS \
            --prune_method $METHOD \
            $SAVE_DIR_PARAM &
    done
    
    # Wait for all parallel jobs of current method to finish
    wait
    echo "All $METHOD pruning tasks completed."
    echo ""
done

echo "All pruning methods executed successfully."