#!/bin/bash
# This is a demo script for torch_pruning pruning experiments on CelebA dataset

# Main training script
TRAIN_SCRIPT="python prune_main.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 \
               --layer_wise_imp \
               --batch_size your_batch_size \
               --train_data_ratio your_train_data_ratio \
               --p-lr your_predictor_lr \
               --ft_epochs your_finetune_epochs \
               ...
               "
# --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

# Pruning methods to test (sequential execution)
PRUNE_METHODS=("method1" "method2" "method3")

# 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
PRUNING_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 dynamic save path
        SAVE_PATH=YOUR_DYNAMIC_SAVE_PATH
        
        echo "Starting $METHOD with seed $SEED on GPU $GPU"
        echo "Save path: $SAVE_PATH"
        
        $TRAIN_SCRIPT $COMMON_PARAMS \
            --gpu $GPU \
            --random_seed $SEED \
            --pruning_ratio $PRUNING_RATIO \
            --iterative_steps $ITERATIVE_STEPS \
            --prune_methods $METHOD \
            --save_dir $SAVE_PATH &
    done
    
    # Wait for all parallel jobs of current method to finish
    wait
    echo "All $METHOD pruning tasks completed."
    echo ""
done

echo "================================================================"
echo " All pruning methods executed successfully! "
echo "================================================================"