DOCSEARCHbyAlgolia
Submit Docs

TrojanZoo

repository·main·Indexed 43 days ago

https://github.com/ain-soph/trojanzoo

A unified PyTorch-based platform for the holistic evaluation of neural backdoors in image classification tasks. It provides tools for training models, executing backdoor attacks such as BadNet, IMC, InvisiblePoison, and Refool, and implementing defenses like Neural Cleanse and AdvMind. The platform includes a suite of vision-specific datasets including ImageNet, GTSRB, and VGGface2, and supports custom module registration for datasets, models, attacks, and defenses.

Tokens
9.5K
Snippets
31
Records
89
Agent score
54%

What's inside TrojanZoo

  1. Use attack-agnostic backdoor defenses in TrojanZoo

    main

    TrojanZoo provides several defense mechanisms designed to be 'attack-agnostic,' meaning they are intended to mitigate backdoor/trojan vulnerabilities without requiring specific knowledge of the underlying attack method.

    Available defense classes in trojanvision.defenses include:

    • AdvTrain: Adversarial Training to improve model robustness.
    • FinePruning: Fine-pruning techniques to remove backdoor triggers by pruning neurons.
    • MagNet: A defense mechanism that uses a reconstruction-based approach to filter inputs.
    • RandomizedSmooth: Randomized smoothing to provide provable robustness.
    • Recompress: Model recompression to eliminate backdoor-related weights.

  2. Use folder-based datasets in TrojanVision

    main

    TrojanVision provides several dataset classes within the trojanvision.datasets module that allow you to load image data from specific directory structures. These classes are designed to handle standard computer vision datasets for Trojan attacks and defenses.

    Supported dataset classes include:

    • CUB200 and CUB200_2011: Bird species datasets.
    • GTSRB: German Traffic Sign Recognition Benchmark.
    • ImageNet: Large-scale image database.
    • ISIC2018: Skin lesion images.
    • VGGface2: Face datasets.

  3. Use TrojanZoo attacks for vision models

    main

    TrojanZoo provides a suite of attack implementations for vision models, categorized into adv (adversarial) and backdoor attacks. You can instantiate these attacks using the create function or by directly using specific attack classes like BackdoorAttack, CleanLabelBackdoor, or DynamicBackdoor from the trojanvision.attacks module.

    from trojanvision.attacks import create

# Example of creating an attack via the factory function
attack = create(attack_type='backdoor', ...)

  4. Use TrojanZoo defenses for TrojanVision

    main

    TrojanVision provides several defense mechanisms to detect or mitigate backdoors in vision models. Defenses are organized into three main categories:

    1. Adversarial Defenses (adv): Techniques designed to counter adversarial attacks.
    2. Backdoor Defenses (backdoor): Specialized methods for detecting or neutralizing backdoor triggers.
    3. Core Defense Classes: Direct access to specific defense implementations via the trojanvision.defenses module.

    You can instantiate specific defense strategies using the following core classes:

    • BackdoorDefense: High-level defense logic for backdoor mitigation.
    • InputFiltering: Defenses that operate by cleaning or filtering input data before it reaches the model.
    • TrainingFiltering: Defenses applied during the training process to prevent backdoor injection.
    • ModelInspection: Defenses that analyze the model weights or internal representations to detect anomalies.

  5. Manage models in TrojanVision

    main

    The trojanvision.models module provides the core interface for interacting with computer vision models within the TrojanZoo framework. It allows you to discover available models, create model instances, and manage model-related arguments.

    Key capabilities include:

    • Discovering models: Use output_available_models to see which models are currently supported.
    • Creating models: Use create to instantiate a model object.
    • Model Abstraction: The ImageModel class serves as the primary interface for interacting with vision models.

  6. Use trojanvision.datasets to manage image data

    main

    The trojanvision.datasets module provides abstractions for handling image datasets used in TrojanZoo vision tasks. It primarily uses two classes to represent data: ImageSet and ImageFolder.

    • ImageSet: A general abstraction for a collection of images.
    • ImageFolder: A specific implementation for loading images organized in a directory structure (typically where subfolders represent classes).

    You can also use create() to instantiate datasets and add_argument() for configuration-related tasks.

    from trojanvision.datasets import ImageFolder

# Example usage of ImageFolder
dataset = ImageFolder(path='path/to/your/images')
# Access members/methods as defined by the class

  7. Use normal datasets in TrojanVision

    main

    TrojanVision provides several standard image datasets for use in vision-based Trojan research. These datasets are available via the trojanvision.datasets module. Supported datasets include:

    • MNIST
    • CIFAR10
    • CIFAR100
    • ImageNet16
    • ImageNet32
    • STL10

    Each class provides methods to load and manage the dataset for training or evaluation purposes.

  8. Manage execution workflows with Process and ModelProcess

    main

    TrojanZoo uses Process and ModelProcess classes to manage the lifecycle and execution of different tasks.

    • Process: A general-purpose class for managing standard execution workflows.
    • ModelProcess: A specialized subclass designed specifically for handling processes related to machine learning models.

  10. How to implement and register DIY Modules

    main

    You can extend TrojanZoo by implementing your own datasets, models, attacks, or defenses. Follow these steps:

    1. Subclass Base Classes: Create your class by subclassing the appropriate base class from trojanzoo or trojanvision.

      • Dataset (or ImageSet, ImageFolder for vision tasks)
      • Model (or ImageModel for vision tasks)
      • Attack
      • Defense

    2. Register the Module: Register your new class in the trojanvision dictionary so the CLI can recognize it.

      # Example registration
from trojanvision.attacks import class_dict
from my_module import MyCustomAttack
class_dict['my_custom_attack'] = MyCustomAttack

    3. Add Configuration (Optional): To avoid modifying code, simply add a {attack_name}.yml or {attack_name}.json file to the relevant configuration directory. The library will automatically pick up these settings.

    trojanvision.attacks.class_dict[attack_name]=AttackClass

  11. Use TrojanZoo backdoor attack classes

    main

    TrojanZoo provides several classes for implementing backdoor attacks in vision tasks. These classes allow you to inject malicious triggers into models during training or via latent space manipulation.

    Available attack implementations include:

    • BadNet: Implements the BadNet backdoor attack pattern.
    • TrojanNN: A neural network-based Trojan attack.
    • IMC: Implements attacks based on the IMC (Implicit Model Corruption) concept.
    • LatentBackdoor: Implements backdoor attacks targeting the latent space of a model.
    • TrojanNet: A general Trojan attack implementation for neural networks.

  12. Basic workflow for backdoor defense

    main

    The standard TrojanZoo workflow for performing backdoor defense involves a sequential creation of components using the trojanvision package. The process follows this lifecycle:

    1. Environment Setup: Initialize the environment with configuration parameters.
    2. Dataset Creation: Load or prepare the dataset.
    3. Model Creation: Initialize the model (optionally using a pretrained model).
    4. Trainer Creation: Set up the trainer for model training or fine-tuning.
    5. Mark Creation: Define the trigger or 'mark' used for the attack.
    6. Attack Creation: Configure the attack using the dataset, model, and mark.
    7. Defense Creation: Configure the defense mechanism using the dataset, model, and attack.
    8. Execution: Run the defense detection using the trainer instance.

    All creation methods (create) accept keyword arguments (**kwargs) parsed from the command line to ensure consistent configuration across all components.

    import trojanvision
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # Register arguments for all components
    trojanvision.environ.add_argument(parser)
    trojanvision.datasets.add_argument(parser)
    trojanvision.models.add_argument(parser)
    trojanvision.trainer.add_argument(parser)
    trojanvision.marks.add_argument(parser)
    trojanvision.attacks.add_argument(parser)
    trojanvision.defenses.add_argument(parser)
    
    kwargs = parser.parse_args().__dict__

    # Sequential component creation
    env = trojanvision.environ.create(**kwargs)
    dataset = trojanvision.datasets.create(**kwargs)
    model = trojanvision.models.create(dataset=dataset, **kwargs)
    trainer = trojanvision.trainer.create(dataset=dataset, model=model, **kwargs)
    mark = trojanvision.marks.create(dataset=dataset, **kwargs)
    attack = trojanvision.attacks.create(dataset=dataset, model=model, mark=mark, **kwargs)
    defense = trojanvision.defenses.create(dataset=dataset, model=model, attack=attack, **kwargs)

    # Execute defense
    if env['verbose']:
        trojanvision.summary(env=env, dataset=dataset, model=model, mark=mark, trainer=trainer, attack=attack, defense=defense)
    
    defense.detect(**trainer)