This project contains tools for model security analysis, including serialization attack detection, membership inference attacks, and adversarial example generation.
- Create and activate a Python virtual environment:
python -m venv venv
source venv/bin/activate # Linux/Mac
# or
venv\Scripts\activate # Windowscd ModelScanning- Install dependencies:
pip install -r requirements.txt-
serialisation/- For detecting malicious model serializationscanner.py- Model file scannerinject.py- Pickle injection utilitieshelper.ipynb- Helper notebook for model downloading and injection
-
mia/- Membership Inference Attack implementationcifar10.py- MIA example on CIFAR10utils.py- Shadow model utilities
-
adversarial/- Adversarial example generation and detectionadv_pgd.py- PGD attack implementationimagenet_classes.txt- ImageNet class labels
You can use the helper.ipynb file to download model and inject a malicious code.
python serialisation/scanner.py /path/to/model.pt
python mia/cifar10.py --target_epochs 12 --attack_epochs 6 --num_shadows 3
cd adversarial
python adv_pgd.py
See requirements.txt for full list of dependencies. Key requirements:
- PyTorch
- TensorFlow
- scikit-learn
- numpy
- matplotlib
- tqdm