#!/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 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 \
               --fairness_eval_flag \
               --fairness_type max_min \
               --layer_wise_imp \
               "
               
# Pruning methods to test (sequential execution)
PRUNE_METHODS=("l1" "FPGM" "method3")

# Seed and GPU mapping (parallel execution within each method)
SEEDS=(2021 2022 2023 2024)
GPUS=(0 1 2 3)

# Pruning configuration
RATIO=0.2  # Set your pruning ratio here
ITERATIVE_STEPS=1  # 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_DYNAMIC_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."