How it Works
TrojanZoo is a PyTorch-based platform designed to facilitate security research in deep learning, with a focus on backdoor attacks and defenses in image classification tasks. It is structured into two primary packages: trojanzoo and trojanvision. The trojanzoo package contains abstract classes and utilities, while trojanvision provides concrete implementations tailored for image classification. This modular design allows researchers to integrate custom datasets, models, attacks, and defenses without modifying the core codebase.
For instance, trojanzoo.datasets.DataSet generates PyTorch-compatible datasets and dataloaders, while trojanvision.datasets.ImageSet and trojanvision.datasets.ImageFolder extend this functionality for image-specific tasks.
TrojanZoo supports Automatic Mixed Precision (AMP) training via PyTorch’s native API, enabling efficient resource utilization during experiments. Configuration files play a central role in defining parameters for datasets, models, attacks, and defenses, with a priority hierarchy that allows user-defined settings to override defaults.
Researchers can register custom modules by subclassing base classes (e.g., Dataset, Model, Attack, or Defense) and integrating them into the library’s configuration system. This approach eliminates the need for code modifications, allowing users to define new attacks or defenses (e.g., IMC or Neural Cleanse) through configuration files alone.
TrojanZoo’s command-line interface (CLI) provides verbose and colorful output options for better visibility during execution. Users can enable these features using flags like --color and --verbose, while --help displays detailed argument usage. The platform also supports tqdm for progress bars, enhancing the user experience during long-running tasks.
Under the hood, TrojanZoo relies on specific dependencies: Python 3.11, PyTorch 2.0.0, and torchvision 0.15.0. While these must be installed manually, the project recommends using conda for dependency management. Installation can be performed via pip, with an emphasis on the editable mode (via pip install -e) to simplify updates. Alternatively, users can deploy TrojanZoo using Docker, with pre-built images available on Docker Hub and GitHub Container Registry (GHCR).
The platform’s design prioritizes reproducibility and modularity, ensuring that experiments can be easily shared and validated. By abstracting low-level details, TrojanZoo allows researchers to focus on high-level security challenges in deep learning, particularly in the context of backdoor threats.
Use Cases
TrojanZoo is primarily used for evaluating backdoor attacks and defenses in deep learning models, particularly those employed in image classification tasks.
- Backdoor Attack Simulation: Researchers can replicate attacks such as BadNet, IMC (Input-Masking Attack), or AdvMind using predefined scripts. For example, launching a BadNet attack on a ResNet18 model trained on CIFAR-10 involves executing a command-line script with specified parameters, including the attack type, dataset, and model architecture.
- Backdoor Defense Evaluation: TrojanZoo supports defenses like Neural Cleanse, which detects and mitigates backdoor triggers in poisoned models. Users can test defenses against specific attacks (e.g., Neural Cleanse against BadNet) by specifying the defense mechanism in the command-line arguments. The platform’s modular design allows for the integration of custom defenses, provided they adhere to the base class structure.
- Adversarial Intent Inference: The platform includes implementations from the AdvMind paper, which focuses on inferring adversary intent in black-box attacks.
- Custom Experimentation: Researchers can develop their own datasets, models, attacks, or defenses by subclassing TrojanZoo’s base classes. The platform provides templates for ImageSet, ImageFolder, and ImageModel, simplifying the integration of new components. Custom modules are registered via configuration files, allowing users to extend functionality without altering the core codebase.
- Reproducibility of Published Work: TrojanZoo includes example scripts that replicate results from associated research papers, such as those presented at EuroS&P 2022, CCS 2020, and KDD 2020. For example, training a ResNet18 model on CIFAR-10 with 95% accuracy can be achieved using the provided train.py script, with additional flags for enabling features like cutout augmentation and gradient clipping.
- Data Loss Prevention for Prompts and Outputs: While primarily focused on backdoor research, TrojanZoo’s modularity can be extended to explore broader security concerns, such as data loss prevention in model inputs and outputs.
The platform’s support for PyTorch’s native APIs and AMP training further enhances its applicability in resource-constrained environments, making it suitable for both academic research and industry-scale evaluations.
Key Features & Capabilities
TrojanZoo’s design emphasizes modularity, flexibility, and reproducibility, offering a suite of features that cater to the needs of security researchers in deep learning.
- Modular and Plug-and-Play Architecture: The platform is divided into two packages: trojanzoo (abstract utilities) and trojanvision (concrete implementations for image classification). For example, a researcher can use trojanzoo’s DataSet class to generate PyTorch-compatible datasets while employing trojanvision’s ImageFolder for image-specific preprocessing.
- Native PyTorch Integration: TrojanZoo’s APIs generate raw PyTorch instances, ensuring compatibility with existing workflows. Datasets produced by trojanzoo.datasets.DataSet are instances of torch.utils.data.Dataset and torch.utils.data.DataLoader, while models inherit from torch.nn.Module and torch.nn.DataParallel. This integration simplifies the adoption of third-party libraries and tools that rely on PyTorch’s ecosystem.
- Automatic Mixed Precision (AMP) Support: The platform enables AMP training via the --amp flag, which optimizes computational efficiency by dynamically adjusting precision during training.
- Flexible Configuration System: Configuration files define parameters for datasets, models, attacks, and defenses, with a priority hierarchy that ensures user-defined settings take precedence. The system supports YAML and JSON formats, allowing researchers to define custom experiments without modifying the codebase.
- Verbose and Colorful Output: TrojanZoo provides detailed logging and visualization options. Users can enable colorful output with --color and verbose mode with --verbose, improving readability during execution.
- Custom Module Registration: Researchers can extend TrojanZoo’s functionality by registering their own datasets, models, attacks, or defenses. This involves subclassing the appropriate base classes (e.g., Dataset, Model, Attack, or Defense) and integrating them into the library’s configuration system. For example, a custom attack like Dynamic Backdoor Injection can be added by defining a new class and updating the trojanvision.attacks.classdict registry.
- For instance, the train.py script trains a ResNet18 model on CIFAR-10 with configurable hyperparameters, while backdoorattack.py and backdoordefense.py demonstrate attack and defense workflows, respectively.
- Docker and Editable Installation: TrojanZoo supports multiple installation methods, including pip, editable mode, and Docker. The editable mode (via pip install -e) allows users to keep the repository updated with git pull, while Docker images are available on Docker Hub and GHCR for containerized deployments.
- Users can navigate the documentation to understand configuration structures, API references, and best practices. The help system also lists all available command-line arguments, reducing the learning curve for new users.
- Support for Advanced Attack and Defense Mechanisms: TrojanZoo implements a variety of backdoor attacks and defenses, including BadNet, IMC, Neural Cleanse, and AdvMind. These implementations are designed to be adaptable, allowing researchers to test hypotheses under different scenarios. For example, AdvMind supports adaptive attacks and defenses, enabling dynamic interactions between adversarial and defensive strategies.
- Data and Model Storage Management: The platform allows users to specify storage paths for datasets, models, attacks, and defenses via configuration files or command-line arguments. For instance, dataset paths are controlled by --datadir, while model outputs are saved to --modeldir, promoting consistency in research workflows.
TrojanZoo’s features collectively address the need for a unified, holistic, and practical evaluation framework for neural backdoors, as outlined in its associated research. By abstracting common functionalities and providing extensible components, the platform lowers the barrier to entry for security research in deep learning.