Merge branch 'Koukyosyumei:main' into main

README.md

@@ -29,7 +29,7 @@
 
 # What is AIJack?
 
-AIJack is an easy-to-use open-source simulation tool for testing the security of your AI system against hijackers. It provides advanced security techniques like *Differential Privacy*, *Homomorphic Encryption*, *K-anonymity* and *Federated Learning* to guarantee protection for your AI. With AIJack, you can test and simulate defenses against various attacks such as *Poisoning*, *Model Inversion*, *Backdoor*, and *Free-Rider*. We support more than 30 state-of-the-art methods. For more information, check our [documentation](https://koukyosyumei.github.io/AIJack/) and start securing your AI today with AIJack.
+AIJack is an easy-to-use open-source simulation tool for testing the security of your AI system against hijackers. It provides advanced security techniques like *Differential Privacy*, *Homomorphic Encryption*, *K-anonymity* and *Federated Learning* to guarantee protection for your AI. With AIJack, you can test and simulate defenses against various attacks such as *Poisoning*, *Model Inversion*, *Backdoor*, and *Free-Rider*. We support more than 30 state-of-the-art methods. For more information, check our [paper](https://arxiv.org/abs/2312.17667) and [documentation](https://koukyosyumei.github.io/AIJack/) and start securing your AI today with AIJack.
 
 # Installation
 
@@ -237,6 +237,10 @@ Below you can find a list of papers and books that either use or extend AIJack.
 - Huang, Shiyuan. A General Framework for Model Adaptation to Meet Practical Constraints in Computer Vision. Diss. Columbia University, 2024.
 - Liu, Can, Jin Wang, and Dongyang Yu. "RAF-GI: Towards Robust, Accurate and Fast-Convergent Gradient Inversion Attack in Federated Learning." arXiv preprint arXiv:2403.08383 (2024).
 
+# Contribution
+
+AIJack welcomes contributions of any kind. If you'd like to address a bug or propose a new feature, please refer to [our guide](docs/source/contribution.rst).
+
 # Contact
 
 welcome2aijack[@]gmail.com

docs/source/aijack.collaborative.fedexp.rst

@@ -0,0 +1,21 @@
+aijack.collaborative.fedexp package
+===================================
+
+Submodules
+----------
+
+aijack.collaborative.fedexp.server module
+-----------------------------------------
+
+.. automodule:: aijack.collaborative.fedexp.server
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: aijack.collaborative.fedexp
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/source/aijack.collaborative.moon.rst

@@ -0,0 +1,21 @@
+aijack.collaborative.moon package
+=================================
+
+Submodules
+----------
+
+aijack.collaborative.moon.client module
+---------------------------------------
+
+.. automodule:: aijack.collaborative.moon.client
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: aijack.collaborative.moon
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/source/aijack.collaborative.rst

@@ -10,10 +10,12 @@ Subpackages
    aijack.collaborative.core
    aijack.collaborative.dsfl
    aijack.collaborative.fedavg
+   aijack.collaborative.fedexp
    aijack.collaborative.fedgems
    aijack.collaborative.fedkd
    aijack.collaborative.fedmd
    aijack.collaborative.fedprox
+   aijack.collaborative.moon
    aijack.collaborative.optimizer
    aijack.collaborative.splitnn
    aijack.collaborative.tree

docs/source/conf.py

@@ -81,4 +81,4 @@
     "style.css",
 ]
 
-jupyter_execute_notebooks = "off"
+nb_execution_mode = "off"

docs/source/contribution.rst

@@ -3,4 +3,153 @@
 .. _contribution:
 
 Contribution Guide
-==================
+==================
+
+Welcome to AIJack's Contribution Guide!
+
+We're thrilled you're interested in contributing to AIJack. This guide outlines the process for submitting code changes and ensuring they adhere to our project's style and formatting conventions.
+
+Getting Started
+---------------
+
+Fork the Repository
+^^^^^^^^^^^^^^^^^^^
+
+* Head over to the AIJack repository on GitHub (https://github.com/Koukyosyumei/AIJack).
+* Click the "Fork" button to create your own copy of the repository.
+
+Clone Your Fork
+^^^^^^^^^^^^^^^
+
+* Open your terminal and navigate to your desired local directory.
+* Use the git clone command to clone your forked repository:
+
+.. code-block:: bash
+
+    git clone https://github.com/<your-username>/AIJack.git
+    # Replace <your-username> with your GitHub username and <project-name> with the actual project name.
+
+Set Up a Development Environment
+--------------------------------
+
+* Build and install AIJack from source code
+
+.. code-block:: bash
+
+    cd AIJack
+
+    # install the dependencies
+    apt install -y libboost-all-dev
+    pip install -U pip
+    pip install "pybind11[global]"
+
+    # install the editable version
+    pip install -e .
+
+
+Coding Style and Formatting
+---------------------------
+
+Google-Style Docstrings
+^^^^^^^^^^^^^^^^^^^^^^^
+
+We use Google-style docstrings to provide clear and consistent documentation for functions, classes, and modules.
+Refer to the Google Python Style Guide (https://github.com/google/styleguide/blob/gh-pages/pyguide.md) for detailed formatting instructions.
+
+Black Code Formatter
+^^^^^^^^^^^^^^^^^^^^
+
+We utilize Black, a popular code formatter, to maintain consistent code style throughout the project.
+
+Ensure Black is installed (pip install black) in your virtual environment.
+
+To format your code before committing, run:
+
+.. code-block:: bash
+
+    black .
+
+
+Isort Import Organizer
+^^^^^^^^^^^^^^^^^^^^^^
+
+isort helps organize imports in a consistent manner.
+
+Install isort (pip install isort) in your virtual environment.
+
+To organize imports, run:
+
+.. code-block:: bash
+
+    isort .
+
+Making Changes
+--------------
+
+Create a Branch
+^^^^^^^^^^^^^^^
+
+* Use git checkout -b <branch-name> to create a new branch for your changes. Replace <branch-name> with a descriptive name (e.g., fix-issue-123).
+
+Implement Your Changes
+^^^^^^^^^^^^^^^^^^^^^^
+
+* Make your code modifications in the appropriate files.
+* Adhere to the coding style and formatting conventions outlined above.
+
+Test Your Changes
+^^^^^^^^^^^^^^^^^
+
+* Write unit tests (if applicable) to verify your code's functionality and prevent regressions.
+
+* Run existing tests with pytest to ensure they still pass after your modifications.
+
+Commit Your Changes
+^^^^^^^^^^^^^^^^^^^
+
+* Stage your changes using
+
+.. code-block:: bash
+
+	git add <file1> <file2>....
+
+* Commit your staged changes with a descriptive message using
+
+.. code-block:: bash
+
+	git commit -m "<commit message>".
+
+Push Your Changes to Your Fork
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* Push your branch to your forked repository on GitHub:
+
+.. code-block:: bash
+
+	git push origin <branch-name>
+
+Submitting a Pull Request
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* Navigate to Your Fork on GitHub:
+* Go to your GitHub repository.
+* Create a Pull Request:
+
+  * Click on the "Pull requests" tab.
+  * Click on "New pull request" and select the branch containing your changes.
+  * Provide a clear and concise title and description for your pull request.
+  * Click on "Create pull request" to submit it for review.
+  * Code Review and Merging
+
+Project maintainers will review your pull request and provide feedback.
+Address any comments or suggestions raised during the review process.
+Once your pull request is approved, it will be merged into the main project repository.
+
+Additional Tips
+---------------
+
+* Consider running black . and isort . before committing your changes to ensure consistent formatting.
+* Provide clear and concise commit messages that describe the purpose of your changes.
+* If you're unsure about anything, feel free to ask for help! You can create an issue on the project's GitHub repository.
+
+Thank you for your contribution to AIJack!

src/aijack/attack/backdoor/dba.py

@@ -39,7 +39,9 @@ def upload_gradients(self):
         def local_train(
             self, local_epoch, criterion, trainloader, optimizer, communication_id=0
         ):
-            for i in range(local_epoch):
+            loss_log = []
+
+            for _ in range(local_epoch):
                 running_loss = 0.0
                 running_data_num = 0
                 for _, data in enumerate(trainloader, 0):
@@ -64,10 +66,9 @@ def local_train(
                     running_loss += loss.item()
                     running_data_num += inputs.shape[0]
 
-                print(
-                    f"communication {communication_id}, epoch {i}: client-{self.user_id+1}",
-                    running_loss / running_data_num,
-                )
+                loss_log.append(running_loss / running_data_num)
+
+            return loss_log
 
     return DistributedBackdoorAttackClientWrapper
 

src/aijack/attack/backdoor/modelreplacement.py

@@ -42,7 +42,9 @@ def upload_gradients(self):
         def local_train(
             self, local_epoch, criterion, trainloader, optimizer, communication_id=0
         ):
-            for i in range(local_epoch):
+            loss_log = []
+
+            for _ in range(local_epoch):
                 if reference_dataloader is not None:
                     running_loss = 0.0
                     running_data_num = 0
@@ -80,10 +82,9 @@ def local_train(
                     running_loss += loss.item()
                     running_data_num += inputs.shape[0]
 
-                print(
-                    f"communication {communication_id}, epoch {i}: client-{self.user_id+1}",
-                    running_loss / running_data_num,
-                )
+                loss_log.append(running_loss / running_data_num)
+
+            return loss_log
 
     return ModelReplacementAttackClientWrapper
 

src/aijack/attack/evasion/diva.py

@@ -4,6 +4,30 @@
 
 
 class DIVAWhiteBoxAttacker(BaseAttacker):
+    """Class implementing the DIVA white-box attack.
+
+    This class provides functionality to perform the DIVA white-box attack on a target model.
+
+    Args:
+        target_model (torch.nn.Module): The target model to be attacked.
+        target_model_on_edge (torch.nn.Module): The target model deployed on the edge.
+        c (float, optional): The trade-off parameter between origin and edge predictions. Defaults to 1.0.
+        num_itr (int, optional): The number of iterations for the attack. Defaults to 1000.
+        eps (float, optional): The maximum perturbation allowed. Defaults to 0.1.
+        lam (float, optional): The step size for gradient updates. Defaults to 0.01.
+        device (str, optional): The device to perform computation on. Defaults to "cpu".
+
+    Attributes:
+        target_model (torch.nn.Module): The target model to be attacked.
+        target_model_on_edge (torch.nn.Module): The target model deployed on the edge.
+        c (float): The trade-off parameter between origin and edge predictions.
+        num_itr (int): The number of iterations for the attack.
+        eps (float): The maximum perturbation allowed.
+        lam (float): The step size for gradient updates.
+        device (str): The device to perform computation on.
+
+    """
+
     def __init__(
         self,
         target_model,
@@ -23,6 +47,16 @@ def __init__(
         self.device = device
 
     def attack(self, data):
+        """Performs the DIVA white-box attack on input data.
+
+        Args:
+            data (tuple): A tuple containing input data and corresponding labels.
+
+        Returns:
+            tuple: A tuple containing the adversarial examples and attack logs.
+
+        """
+
         x, y = data
         x = x.to(self.device)
         y = y.to(self.device)