doc/tutorial/fed_ipr/FedIPR-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FedIPR Tutorial: Guide to Adding Watermarks to Image and Language Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, you'll learn how to add both backdoor-based and feature-based watermarks to your models in the federated training. \n",
    "We'll dive into using backdoor-watermark datasets for backdoor-based watermarking and exploring signblock—a tool that learns feature-based watermarks during traning. We will show you how to apply these techniques to both computer vision and language models. We'll also offer a hands-on example with a CV task, share how to verify the watermarks you've embedded, and introduce some ready-to-use models provided by the FATE framework. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FedIPR Introduction\n",
    "FedIPR stands for Federated Intellectual Property Rights, a technology designed to protect the ownership of models developed under federated conditions. At its core, the FedIPR approach is described in the original paper [FedIPR](https://arxiv.org/pdf/2109.13236.pdf), introducing two primary watermarking techniques to safeguard your model: Backdoor-based and Feature-based watermarks.\n",
    "\n",
    "Backdoor-based methods: These methods use specific input triggers to produce intentional, incorrect labels. The goal here is to create a unique \"signature\" for the model, allowing for ownership verification through remote APIs, without requiring access to the model's internal parameters.\n",
    "\n",
    "Feature-based methods: These techniques encode designated binary strings as watermarks directly into the model's layer parameters. Various schemes have been proposed, such as embedding these watermarks into convolution layer weights using a binary cross-entropy loss function, or into normalization layer scale parameters using a hinge-like regularization term. In our implementations, we embed signatures into normalization layers as the same as \n",
    "\n",
    "Through these watermarking techniques, FedIPR ensures a robust way to assert ownership of your federated models without compromising their performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preliminary\n",
    "\n",
    "We strongly recommend you finish reading our NN tutorial to get familiar with Model and Dataset customizations: [NN Tutorials](https://github.com/FederatedAI/FATE/blob/master/doc/tutorial/pipeline/nn_tutorial/README.md)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backdoor Dataset for Backdoor Watermark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can develop your own backdoor dataset and use it in FedIPRTrainer. If watermark dataset is detected, it will be used to train models along with your task dataset. If not provided, it will perform normal training.\n",
    "\n",
    "You can add python path so that you can run codes in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "your_path_to_fate_python = 'xxx/fate/python'\n",
    "sys.path.append(your_path_to_fate_python)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interfaces"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The WaterMarkDataset class serves as a base class for handling watermark datasets in federated learning environments. It’s crucial for you to implement the load method. The primary task when subclassing WaterMarkDataset is to fill in the load method. This method should take a path argument and use it to load both your normal and watermark datasets.\n",
    "\n",
    "Besides you need to implement other interfaces like get_item, len like using a pytorch dataset to make it work correctly in FATE.\n",
    "You can refer to this tutorial: [Dataset Customization](https://github.com/FederatedAI/FATE/blob/master/doc/tutorial/pipeline/nn_tutorial/Homo-NN-Customize-your-Dataset.ipynb)\n",
    "\n",
    "Here show you the source code of the watermark dataset class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from federatedml.nn.dataset.base import Dataset\n",
    "from federatedml.util import LOGGER\n",
    "from federatedml.nn.dataset.image import ImageDataset\n",
    "\n",
    "\n",
    "class WaterMarkDataset(Dataset):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.normal_dataset = None\n",
    "        self.watermark_dataset = None\n",
    "\n",
    "    def load(self, path):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def get_normal_dataset(self):\n",
    "        return self.normal_dataset\n",
    "\n",
    "    def get_watermark_dataset(self):\n",
    "        return self.watermark_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make you better understand how our watermark dataset work, here we show the implementation of load function of our built-in WaterMarkImageDataset.\n",
    "The WaterMarkImageDataset class is designed to automatically identify and load two distinct folders from the specified file path: one containing 'normal' training samples and another containing 'watermark' trigger samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load(self, file_path):\n",
    "\n",
    "    # normal dataset path\n",
    "    normal_path = os.path.join(file_path, self.normal_folder_name)\n",
    "    # watermark dataset path\n",
    "    watermark_path = os.path.join(file_path, self.watermark_folder_name)\n",
    "\n",
    "    # load normal dataset\n",
    "    self.normal_dataset = ImageDataset(\n",
    "        center_crop=self.center_crop,\n",
    "        center_crop_shape=self.size,\n",
    "        generate_id_from_file_name=self.generate_id_from_file_name,\n",
    "        file_suffix=self.file_suffix,\n",
    "        float64=self.float64,\n",
    "        label_dtype=self.label_type\n",
    "    )\n",
    "    if os.path.exists(normal_path):\n",
    "        self.normal_dataset.load(normal_path)\n",
    "    else:\n",
    "        self.normal_dataset = None\n",
    "        LOGGER.info(\n",
    "            f'normal dataset not found in {normal_path}, will not load normal dataset')\n",
    "    # load watermark dataset\n",
    "    self.watermark_dataset = ImageDataset(\n",
    "        center_crop=self.center_crop,\n",
    "        center_crop_shape=self.size,\n",
    "        generate_id_from_file_name=self.generate_id_from_file_name,\n",
    "        file_suffix=self.file_suffix,\n",
    "        float64=self.float64,\n",
    "        label_dtype=self.label_type\n",
    "    )\n",
    "    if os.path.exists(watermark_path):\n",
    "        self.watermark_dataset.load(watermark_path)\n",
    "    else:\n",
    "        self.watermark_dataset = None\n",
    "        LOGGER.info(\n",
    "            f'watermark dataset not found in {watermark_path}, will not load watermark dataset')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can try our WaterMarkImageDataset: use it load our provided cifar-10 watermarked dataset which contains 100 trigger samples.Each image in these folders has been augmented with a pattern of structured noise in one corner. Download the dataset and place it in example/data folder in your fate project: [Dowload Path]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.dataset.watermark import WaterMarkImageDataset\n",
    "\n",
    "ds = WaterMarkImageDataset()\n",
    "ds.load('../../../examples/data/cifar_10_ipr/fedipr_cifar10_guest/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset ImageFolder\n",
       "    Number of datapoints: 25000\n",
       "    Root location: ../../examples/data/cifar_10_ipr/fedipr_cifar10_guest/normal\n",
       "    StandardTransform\n",
       "Transform: Compose(\n",
       "               ToTensor()\n",
       "           )"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.get_normal_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset ImageFolder\n",
       "    Number of datapoints: 100\n",
       "    Root location: ../../examples/data/cifar_10_ipr/fedipr_cifar10_guest/watermark\n",
       "    StandardTransform\n",
       "Transform: Compose(\n",
       "               ToTensor()\n",
       "           )"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.get_watermark_dataset() # water mark dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25100"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, you can now customize a watermark dataset for your own tasks to add watermarks to your models. In the upcoming CIFAR-10 task, we will be using FATE's built-in image watermark dataset class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Built-in BacthNorm and LayerNorm Blocks for Feature-based Watermark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will delve into the workings of feature-based watermarking. Feature-based watermarking involves embedding binary watermarks vectors into specific model parameters. In FATE implementations, we use the same design as the FATE-IPR paper: In the case of CNN, binary water mark are embeded into BatchNorm Layer. In transformers, watermarks are embeded into LayerNorm layers.\n",
    "\n",
    "You can use SignatureConv, SignatureLayerNorm to build your model. Once these blocks are detected in the FedIPR trainer, trainer will automatically assign binary watermark vector whose bit length is computed by Equation (15) in the origin paper.\n",
    "\n",
    "You can import them from:model's proprietary elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.model_zoo.ipr.sign_block import SignatureConv, SignatureLayerNorm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we show you the source code of our built in alexnet and distilbert to show you how to quickly build a model with featurebased watermark:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from fate_llm.model_zoo.ipr.sign_block import SignatureConv, ConvBlock\n",
    "\n",
    "\n",
    "class SignAlexNet(nn.Module):\n",
    "\n",
    "    \"\"\"\n",
    "    This is a modified Alexnet: its 4,5,6 layers are replaced by Singnature Conv Block\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_classes):\n",
    "        super().__init__()\n",
    "        in_channels = 3\n",
    "        maxpoolidx = [1, 3, 7]\n",
    "        signed_layer = [4, 5, 6]\n",
    "        layers = []\n",
    "        inp = in_channels\n",
    "\n",
    "        # channels & kennel size\n",
    "        # the same setting as the FedIPR paper\n",
    "        oups = {\n",
    "            0: 64,\n",
    "            2: 192,\n",
    "            4: 384,\n",
    "            5: 256,\n",
    "            6: 256\n",
    "        }\n",
    "        kp = {\n",
    "            0: (5, 2),\n",
    "            2: (5, 2),\n",
    "            4: (3, 1),\n",
    "            5: (3, 1),\n",
    "            6: (3, 1)\n",
    "        }\n",
    "\n",
    "        for layeridx in range(8):\n",
    "            if layeridx in maxpoolidx:\n",
    "                layers.append(nn.MaxPool2d(2, 2))\n",
    "            else:\n",
    "                k = kp[layeridx][0]\n",
    "                p = kp[layeridx][1]\n",
    "                if layeridx in signed_layer:\n",
    "                    layers.append(SignatureConv(inp, oups[layeridx], k, 1, p))\n",
    "                else:\n",
    "                    layers.append(ConvBlock(inp, oups[layeridx], k, 1, p))\n",
    "                inp = oups[layeridx]\n",
    "\n",
    "        self.features = nn.Sequential(*layers)\n",
    "        self.classifier = nn.Linear(4 * 4 * 256, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        for m in self.features:\n",
    "            x = m(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.classifier(x)\n",
    "        if self.training:\n",
    "            return x\n",
    "        else:  # Sofmax\n",
    "            return nn.functional.softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By inserting signconv block you can easily build a cv model with feature-based signature, in the case of NLP models, by useing 'recursive_replace_layernorm' you can quickly replace the original LayerNorm with our sign layernorm. Codes below show that you can quickly add feature-based watermarks to a huggingface pretraind model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Module\n",
    "from transformers import DistilBertForSequenceClassification, DistilBertForTokenClassification\n",
    "from fate_llm.model_zoo.ipr.sign_block import recursive_replace_layernorm\n",
    "\n",
    "\n",
    "class SignDistilBertForTokenClassification(Module):\n",
    "\n",
    "    def __init__(self, model_path=None, num_labels=4) -> None:\n",
    "        super().__init__()\n",
    "        if model_path is None:\n",
    "            model_path = 'distilbert-base-uncased'\n",
    "\n",
    "        self.model_path = model_path\n",
    "        self.model = DistilBertForTokenClassification.from_pretrained(\n",
    "            model_path, num_labels=num_labels)\n",
    "\n",
    "        # replace layernorm by SignatureLayerNorm\n",
    "        sub_distilbert = self.model.distilbert.transformer.layer[3:]  # replace layernorm by SingLayerNorm in the last 3 layer\n",
    "        recursive_replace_layernorm(\n",
    "            sub_distilbert,\n",
    "            layer_name_set={'output_layer_norm'})\n",
    "\n",
    "    def forward(self, input_dict):\n",
    "        return self.model(**input_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verify Feature-based watermark with our tools\n",
    "\n",
    "After training is done, feature-based watermarks' signatures will be saved together with model. You can use our tool to verify the model ownership."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.trainer.fedipr_trainer import verify_feature_based_signature"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the example below for usage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FedIPR on FATE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In FATE-LLM-1.3’s model_zoo we have these built-in models which are automatically integrated with feature-based watermarking capabilities:\n",
    "\n",
    "#### Model List\n",
    "\n",
    "- `alexnet.py` - Alexnet\n",
    "- `resnet.py` - Resnet18\n",
    "- `distilbert.py` - Distilbert (Configurations match those in the FedIPR paper)\n",
    "- `gpt2.py` - Standard GPT-2 (Watermarks are added to the last 2 transformer layers)\n",
    "t.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have verified the effectiveness of our watermarking features through a series of tests:\n",
    "- For computer vision tasks, we evaluated both backdoor watermarking and feature-based watermarking techniques on the CIFAR-10 and CIFAR-100 datasets. Our testing involved the use of ResNet and AlexNet models.\n",
    "- For natural language processing tasks, we assessed the performance of DistilBERT and GPT2 models on the IMDB & CoNLL-2003 datasets, which are sequence classification tasn and token classification respectively. \n",
    "During the testing phase, the sign bit was automatically allocated, and the data was evenly divided between the guest and host parties. For backdoor watermarking, each party supplied 100 trigger samples, all of which were augmented with noises.\n",
    "\n",
    "Here we display the results of the experiments:\n",
    "\n",
    "AlexNet & Resnet:\n",
    "\n",
    "| Test Configuration | AlexNet Feature-Based Watermark Accuracy | AlexNet Backdoor Watermark Accuracy | ResNet18 Feature-Based Watermark Accuracy | ResNet18 Backdoor Watermark Accuracy |\n",
    "|--------------------|-----------------------------------------|------------------------------------|------------------------------------------|-------------------------------------|\n",
    "| Two-party federation on CIFAR-10 with 100 trigger samples, SignBit auto-assigned | 1.0 (All Parties) | 1.0 (All Parties) | 1.0 (All Parties) | 1.0 (All Parties) |\n",
    "| Two-party federation on CIFAR-100 with 100 trigger samples, SignBit auto-assigned | 1.0 (All Parties) | 1.0 (Guest), 0.991 (Host) | 1.0 (All Parties) | 1.0 (All Parties) |\n",
    "\n",
    "DistilBert & GPT2:\n",
    "\n",
    "| Test Configuration | DistillBERT Feature-Based Watermark Accuracy | GPT-2 Feature-Based Watermark Accuracy |\n",
    "|--------------------|----------------------------------------------|---------------------------------------|\n",
    "| Two-party federation on CoNLL-2003 Token Classification with SignBit auto-assigned | 1.0 (All Parties) | 1.0 (All Parties) |\n",
    "| Two-party federation on IMDB Classification with SignBit auto-assigned | 1.0 (All Parties) | 1.0 (All Parties) |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Cifar-10 Example &  Verifying Watermark\n",
    "\n",
    "At last, we will show you a CV example: we will train a AlexNet with backdoor watermark & feature-based watermark at the same time. And after training is done, we use built in tools to verify feature-based watermark. You can verify the backdoor watermark yourself by simply predicting trigger samples with your models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FedIPR Parameters\n",
    "\n",
    "The FedIPRTrainer's parameters are basically the same as the FedAVGTrainer except for 3 parameters: alpha, verify_freq and backdoor_verify_method\n",
    "alpha is the weight for sign loss; verify_freq is the frequency of verifying your watermark during training(you can check result in logs) and backdoor_verify_method allows you to choose the method for verifying your datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FedIPRTrainer(FedAVGTrainer):\n",
    "\n",
    "    def __init__(self,\n",
    "                 epochs=10,\n",
    "                 noraml_dataset_batch_size=32,\n",
    "                 watermark_dataset_batch_size=2,\n",
    "                 early_stop=None,\n",
    "                 tol=0.0001,\n",
    "                 secure_aggregate=True,\n",
    "                 weighted_aggregation=True,\n",
    "                 aggregate_every_n_epoch=None,\n",
    "                 cuda=None,\n",
    "                 pin_memory=True,\n",
    "                 shuffle=True,\n",
    "                 data_loader_worker=0,\n",
    "                 validation_freqs=None,\n",
    "                 checkpoint_save_freqs=None,\n",
    "                 task_type='auto',\n",
    "                 save_to_local_dir=False,\n",
    "                 collate_fn=None,\n",
    "                 collate_fn_params=None,\n",
    "                 alpha=0.01,\n",
    "                 verify_freqs=1,\n",
    "                 backdoor_verify_method: Literal['accuracy',\n",
    "                                                 'loss'] = 'accuracy'):\n",
    "        ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Submit a pipeline to run FedIPR CV task\n",
    "\n",
    "This a standalone version example, if you are running on the cluster version, you have to bind name&namespace on guest&host machines correspondingly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pipeline.backend.pipeline.PipeLine at 0x7fc070abdc40>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch as t\n",
    "from torch import nn\n",
    "from pipeline import fate_torch_hook\n",
    "from pipeline.component import HomoNN\n",
    "from pipeline.backend.pipeline import PipeLine\n",
    "from pipeline.component import Reader, Evaluation, DataTransform\n",
    "from pipeline.interface import Data, Model\n",
    "\n",
    "t = fate_torch_hook(t)\n",
    "\n",
    "import os\n",
    "# bind data path to name & namespace\n",
    "fate_project_path = os.path.abspath('../../../')\n",
    "host = 9997\n",
    "guest = 9997\n",
    "arbiter = 9997\n",
    "pipeline = PipeLine().set_initiator(role='guest', party_id=guest).set_roles(guest=guest, host=host,\n",
    "                                                                            arbiter=arbiter)\n",
    "\n",
    "data_0 = {\"name\": \"watermark_cifar10_guest\", \"namespace\": \"experiment\"}\n",
    "data_1 = {\"name\": \"watermark_cifar10_host\", \"namespace\": \"experiment\"}\n",
    "\n",
    "data_path_0 = fate_project_path + '/examples/data/cifar_10_ipr/fedipr_cifar10_guest'\n",
    "data_path_1 = fate_project_path + '/examples/data/cifar_10_ipr/fedipr_cifar10_host'\n",
    "pipeline.bind_table(name=data_0['name'], namespace=data_0['namespace'], path=data_path_0)\n",
    "pipeline.bind_table(name=data_1['name'], namespace=data_1['namespace'], path=data_path_1)\n",
    "\n",
    "reader_0 = Reader(name=\"reader_0\")\n",
    "reader_0.get_party_instance(role='guest', party_id=guest).component_param(table=data_0)\n",
    "reader_0.get_party_instance(role='host', party_id=host).component_param(table=data_1)\n",
    "\n",
    "from pipeline.component.nn import DatasetParam\n",
    "\n",
    "dataset_param = DatasetParam(dataset_name='watermark')\n",
    "\n",
    "from pipeline.component.homo_nn import TrainerParam  # Interface\n",
    "\n",
    "# our simple classification model:\n",
    "model = t.nn.CustModel(module_name='ipr.alexnet', class_name='SignAlexNet', num_classes=10)\n",
    "\n",
    "nn_component = HomoNN(name='nn_0',\n",
    "                      model=model, # model\n",
    "                      dataset=dataset_param,  # dataset\n",
    "                      # Notice that for the convenience of getting result model we set save_to_local_dir=True\n",
    "                      trainer=TrainerParam(trainer_name='fedipr_trainer', epochs=5, save_to_local_dir=True, cuda=0),\n",
    "                      optimizer=t.optim.Adam(lr=0.001),\n",
    "                      loss=t.nn.CrossEntropyLoss(),\n",
    "                      torch_seed=100 # random seed\n",
    "                      )\n",
    "\n",
    "\n",
    "pipeline.add_component(reader_0)\n",
    "pipeline.add_component(nn_component, data=Data(train_data=reader_0.output.data))\n",
    "pipeline.compile()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.fit() # submit!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Model and Verify\n",
    "\n",
    "Since we enable 'save_to_local_dir', we can directly load trained model from fateflow job folder, and verify its watermarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.trainer.fedipr_trainer import verify_feature_based_signature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_id = '202309041103336933850'  # your job id\n",
    "import os\n",
    "fate_project_path = os.path.abspath('../../../')\n",
    "local_dir = fate_project_path + '/fateflow/jobs/{}/guest/9997/nn_0/'.format(job_id)\n",
    "state_dict = t.load(local_dir + 'model.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.model_zoo.ipr.alexnet import SignAlexNet\n",
    "\n",
    "model = SignAlexNet(num_classes=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(state_dict['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = state_dict['extra_data']['keys']  # W and watermark vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'features.4': (tensor([[-7.3380e-02,  1.6275e+00, -1.5404e+00,  ...,  3.4250e-01,\n",
       "           -1.0067e+00, -5.4504e-01],\n",
       "          [ 2.9928e-01, -4.0935e-01, -6.1239e-01,  ...,  7.2356e-01,\n",
       "            2.7019e-01, -9.1200e-01],\n",
       "          [-4.3889e-02,  2.1774e+00, -1.3706e+00,  ..., -8.5879e-01,\n",
       "            2.3445e-01,  2.0458e+00],\n",
       "          ...,\n",
       "          [-5.1755e-01,  5.9240e-01,  2.6353e-01,  ..., -1.0465e+00,\n",
       "           -5.3456e-01, -6.0439e-01],\n",
       "          [-2.4679e-01, -1.4290e+00, -5.9567e-01,  ...,  7.7682e-01,\n",
       "           -6.2445e-01,  1.3682e+00],\n",
       "          [ 1.1148e+00, -8.7518e-01,  7.6818e-01,  ...,  6.5654e-01,\n",
       "           -1.8362e+00, -5.5355e-04]]),\n",
       "  tensor([-1., -1.,  1.,  1.,  1., -1.,  1.,  1., -1., -1., -1.,  1., -1.,  1.,\n",
       "          -1., -1.,  1., -1., -1., -1., -1.,  1., -1.,  1., -1., -1.,  1., -1.,\n",
       "           1., -1., -1., -1.,  1.,  1.,  1.,  1.,  1.,  1.,  1., -1., -1.,  1.,\n",
       "           1., -1.,  1., -1., -1.,  1., -1.,  1., -1., -1., -1.,  1., -1., -1.,\n",
       "          -1., -1., -1., -1.,  1., -1., -1.,  1.,  1., -1., -1.,  1., -1., -1.,\n",
       "           1.,  1.,  1., -1., -1., -1., -1.,  1., -1.,  1., -1.,  1., -1.,  1.,\n",
       "          -1., -1.,  1., -1., -1., -1., -1.,  1., -1.,  1.,  1., -1., -1.,  1.,\n",
       "          -1., -1.,  1., -1., -1., -1., -1.,  1.,  1.,  1., -1., -1.,  1., -1.,\n",
       "          -1., -1.,  1., -1.,  1., -1., -1., -1., -1.,  1., -1.,  1.,  1., -1.,\n",
       "           1.,  1.,  1., -1.,  1.,  1., -1., -1.,  1.,  1., -1., -1., -1.,  1.,\n",
       "          -1., -1., -1., -1.,  1., -1., -1.,  1.,  1.,  1., -1.,  1., -1.,  1.,\n",
       "           1.,  1., -1.,  1.,  1.,  1.,  1.,  1.,  1.,  1., -1., -1., -1.,  1.,\n",
       "          -1., -1., -1., -1., -1., -1., -1., -1., -1.,  1., -1., -1.,  1.,  1.,\n",
       "           1., -1.,  1., -1., -1., -1.,  1.,  1.,  1.])),\n",
       " 'features.5': (tensor([[-1.2336,  0.1894, -0.3584,  ..., -0.5398,  0.5318, -1.6536],\n",
       "          [ 0.1128,  0.3999,  1.2841,  ...,  1.6082, -0.1920, -0.0636],\n",
       "          [-0.9447, -0.2025,  0.4786,  ...,  1.5100, -0.7834,  0.8102],\n",
       "          ...,\n",
       "          [-0.7941,  2.0311, -0.9690,  ..., -1.1630,  0.6953,  1.6115],\n",
       "          [ 0.0314,  0.3718,  0.5974,  ..., -1.6695,  1.8833, -0.1461],\n",
       "          [ 0.4956,  0.7747, -0.0847,  ..., -0.3533,  0.0763,  0.0952]]),\n",
       "  tensor([-1.,  1., -1.,  1.,  1.,  1.,  1.,  1.,  1., -1., -1., -1.,  1., -1.,\n",
       "          -1., -1., -1.,  1., -1.,  1., -1., -1.,  1., -1., -1., -1.,  1., -1.,\n",
       "           1.,  1.,  1., -1., -1., -1.,  1.,  1.,  1.,  1.,  1.,  1., -1.,  1.,\n",
       "          -1., -1., -1., -1.,  1., -1.,  1., -1.,  1., -1.,  1., -1., -1.,  1.,\n",
       "           1.,  1.,  1.,  1., -1.,  1.,  1.,  1., -1., -1., -1., -1., -1., -1.,\n",
       "          -1.,  1.,  1., -1.,  1., -1., -1.,  1., -1., -1.,  1.,  1., -1., -1.,\n",
       "          -1., -1., -1.,  1.,  1.,  1.,  1.,  1.,  1., -1.,  1., -1., -1.,  1.,\n",
       "           1.,  1.,  1., -1.,  1., -1.,  1.,  1., -1., -1.,  1.,  1., -1., -1.,\n",
       "           1., -1., -1.,  1.,  1., -1.,  1.,  1.,  1.,  1., -1., -1., -1., -1.,\n",
       "           1.])),\n",
       " 'features.6': (tensor([[ 2.6993e+00,  1.0507e+00, -6.6219e-01,  ...,  6.3679e-01,\n",
       "            7.7061e-01,  1.4231e+00],\n",
       "          [-1.0477e+00,  2.0904e-01, -3.4522e-01,  ..., -4.9581e-01,\n",
       "            1.4211e+00, -2.1041e+00],\n",
       "          [ 1.0036e+00,  1.0025e+00, -2.5215e-03,  ...,  1.1413e+00,\n",
       "           -1.8600e+00,  2.0058e-02],\n",
       "          ...,\n",
       "          [ 1.2943e+00,  5.6073e-01, -1.9590e+00,  ..., -1.4320e+00,\n",
       "           -1.6486e+00, -3.0871e-01],\n",
       "          [ 4.2747e-01,  1.8310e+00, -2.7685e-01,  ..., -1.0765e+00,\n",
       "           -4.6004e-01,  3.6701e-02],\n",
       "          [-4.9978e-01,  4.4728e-01, -7.3183e-01,  ...,  7.5242e-01,\n",
       "            8.4118e-01,  8.3414e-02]]),\n",
       "  tensor([ 1., -1.,  1.,  1., -1., -1.,  1., -1., -1.,  1.,  1.,  1.,  1., -1.,\n",
       "          -1., -1., -1.,  1.,  1.,  1.,  1., -1., -1., -1.,  1.,  1.,  1.,  1.,\n",
       "          -1.,  1.,  1., -1.,  1., -1., -1.,  1., -1., -1.,  1.,  1., -1.,  1.,\n",
       "          -1., -1., -1., -1., -1.,  1.,  1., -1.,  1., -1., -1., -1., -1., -1.,\n",
       "          -1., -1.,  1., -1.,  1.,  1., -1.,  1., -1., -1.,  1., -1., -1.,  1.,\n",
       "           1.,  1., -1., -1., -1.,  1., -1., -1.,  1., -1.,  1., -1.,  1., -1.,\n",
       "          -1.,  1.,  1.,  1., -1.,  1.,  1.,  1.,  1.,  1.,  1., -1.,  1., -1.,\n",
       "          -1.,  1.,  1.,  1., -1.,  1.,  1., -1., -1.,  1.,  1.,  1.,  1.,  1.,\n",
       "          -1., -1., -1., -1.,  1.,  1.,  1., -1.,  1., -1.,  1., -1.,  1.,  1.,\n",
       "          -1., -1., -1.]))}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n"
     ]
    }
   ],
   "source": [
    "acc = verify_feature_based_signature(model, keys)\n",
    "print(acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The accuracy is 100%! Congratulations. Now you can use FATE to build your own IPR protected models."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}