# Overview

This directory provides a complete pipeline for structured pruning, including pretraining base models, different structured pruning approaches implemented based on torch_pruning library([GitHub - VainF/Torch-Pruning: [CVPR 2023] DepGraph: Towards Any Structural Pruning; LLMs, Vision Foundation Models, etc.](https://github.com/VainF/Torch-Pruning)) and **our proposed method** **F-PVE**.

## Project Structure

```
.
├── archs/                  # Network architectures
│   ├── resnet_imagenet.py # ResNet for CelebA
├── pretrain/              # Pretraining related files
├── prune_tp/             # Torch_pruning implementation
└── utils/                # Utility functions
```

## Workflow

### 1. Pretraining Base Models

Before pruning, you need to pretrain your base model:

```bash
cd pretrain
bash pretrain.sh
```

This will train your selected model architecture and save the model weights for pruning.

### 2. Model Pruning

You can choose between two pruning approaches:

#### Option 1: Torch Pruning (TP)

Execute the TP pruning script:

```bash
cd prune_tp
bash scripts/run_tp.sh
```

This will execute the pruning process using torch_pruning library, followed by model fine-tuning.

#### Option 2: F-PVE (Our Method)

To use our proposed F-PVE pruning method:

```bash
bash FPVE_scripts/run_FPVE.sh
```

## Getting Started

1. First, choose and pretrain your base model using the pretraining script
2. Select your pruning method (methods implemented by torch_pruning or F-PVE)
3. Execute the corresponding pruning script
4. Evaluate the performance of pruned models 

For detailed configuration options and parameters, refer to the respective script files.

## Key Components

- `pretrain/trainer.py`: Base model training implementation
- `prune_tp/prune_main.py`: Main pruning logic using torch_pruning
- `prune_tp/prune_finetune.py`: Fine-tuning implementation for pruned models
- `FPVE.py`: Implementation of our proposed F-PVE pruning method

## Configuration

- Modify `pretrain.sh` to customize your pretraining setup
- Adjust pruning parameters in `run_tp.sh` or `run_FPVE.sh` based on your requirements
- Set the correct path for the model and data
