{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Model Inversion Attack \n\nThe objective of a **Model Inversion Attack** is to recover information about the data the model was trained on, or outright data samples. The model in question is referred to as a **Target Model** and is a so-called *Black-Box Model* since we don't have information about its architecture, its parameters, the way it was trained and the data used. We do however have access to its prediction function; we can access it through an API Query for instance.\n\nBased on the purpose of this target model, we can infer the type of data it was used on. For instance, a model that is used to predict if an image is a dog or a hotdog was trained on samples of dogs and hotdogs. Note that it is certainly not necessary to have the exact same training set. So we can use any creative data collection technique to gather a dataset of a similar distribution to the supposed training set.\n\nNow the pipeline of this attack is the following: first, we collect data samples and then we send requests to the Target Model to record the predictions made. Finally, we train an **Attack Model** to reconstruct images given the output of the target model. In other words, the input to the Attack Model is a Tensor of size *m* (where m is the number of classes/outputs of the Target Model) and its output is another tensor of size (n, n, 3) where is the size of the original image fed to the Target Model. In other words, we build have an Encoder-Decoder architecture, where the encoder is the Target Model and the decoder is our Attack Model.\n\nThe final step of the model inversion attack is the data recovery stage: we query the Attack model with made up tensors. Those tensors can be either a one hot-encoded for a specific class, in which case the Attack Model will return the image that is most likely to be predicted for each of the *m* dimensions of the model.\n\nIn this notebook, we will query an API created for the competition [AI Village Capture the Flag @ DEFCON](https://www.kaggle.com/competitions/ai-village-ctf?rvi=1), specifically the *Inference Challenge*. The objective is to recover a 6 one letter images that activate the classes of the Target Model. Knowing that the training data was letters, we can use a dataset of handwritten letters. The architecture of the pipeline will once again be a Target Model that takes images as input classifies/encodes then in 6 classes/dimensions and an Attack Model that takes the 6 dimensional data and upscales them to the original images.\n\nThis Pytorch Implementation is inspired by [OpenMined](https://github.com/OpenMined/PySyft/blob/a27deed0d07c199de039fafd323164640c9c8f6d/examples/tutorials/advanced/privacy_attacks/Tutorial%201%20-%20Black%20box%20model%20inversion.ipynb).","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport requests\nimport json\nfrom collections import namedtuple\n\nimport matplotlib.pyplot as plt\nfrom matplotlib.image import imread\n\nimport tqdm\nimport torch\nimport torchvision.transforms as transforms\nfrom torchvision.datasets import EMNIST, MNIST\nfrom tqdm.notebook import tqdm, trange\nfrom torch.optim.lr_scheduler import StepLR\n\n! pip install torchsummary \nfrom torchsummary import summary","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:35:30.635614Z","iopub.execute_input":"2022-09-13T21:35:30.636163Z","iopub.status.idle":"2022-09-13T21:35:41.667588Z","shell.execute_reply.started":"2022-09-13T21:35:30.636121Z","shell.execute_reply":"2022-09-13T21:35:41.666394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Processing","metadata":{}},{"cell_type":"code","source":"# Image processing\n# This steps is only needed because the EMNIST Dataset used has flipped images\nmnist_transform = transforms.Compose(\n    [\n        lambda img: transforms.functional.rotate(img, -90),\n        lambda img: transforms.functional.hflip(img),\n        transforms.Resize((32, 32)),\n        transforms.ToTensor(),\n    ]\n)\n\n# Get dataset\nattacker_dataset_train = EMNIST(\"emnist\", 'balanced', download=True, train=True, transform=mnist_transform)\nattacker_dataset_test = EMNIST(\"emnist\", 'balanced', download=True, train=False, transform=mnist_transform)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:09.021678Z","iopub.execute_input":"2022-09-13T21:34:09.022615Z","iopub.status.idle":"2022-09-13T21:34:34.162966Z","shell.execute_reply.started":"2022-09-13T21:34:09.022572Z","shell.execute_reply":"2022-09-13T21:34:34.161787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Remove lower case letters\nattacker_dataset_train.data = attacker_dataset_train.data[attacker_dataset_train.targets <= 35]\nattacker_dataset_train.targets = attacker_dataset_train.targets[attacker_dataset_train.targets <= 35]\n\nattacker_dataset_test.data = attacker_dataset_test.data[attacker_dataset_test.targets <= 35]\nattacker_dataset_test.targets = attacker_dataset_test.targets[attacker_dataset_test.targets <= 35]\n\n# Keep fewer samples (to reduce execution time of label collection)\nattacker_dataset_train.data = attacker_dataset_train.data[:10000]\nattacker_dataset_train.targets = attacker_dataset_train.targets[:10000]\n\nattacker_dataset_test.data = attacker_dataset_test.data[:100]\nattacker_dataset_test.targets = attacker_dataset_test.targets[:100]","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:34.164552Z","iopub.execute_input":"2022-09-13T21:34:34.164906Z","iopub.status.idle":"2022-09-13T21:34:34.534371Z","shell.execute_reply.started":"2022-09-13T21:34:34.164869Z","shell.execute_reply":"2022-09-13T21:34:34.533387Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dataloader\nattacker_loader_train = torch.utils.data.DataLoader(attacker_dataset_train, batch_size=64)\nattacker_loader_test = torch.utils.data.DataLoader(attacker_dataset_test, batch_size=64)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:34.536793Z","iopub.execute_input":"2022-09-13T21:34:34.537161Z","iopub.status.idle":"2022-09-13T21:34:34.546911Z","shell.execute_reply.started":"2022-09-13T21:34:34.537125Z","shell.execute_reply":"2022-09-13T21:34:34.544658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_images(tensors):\n    \"\"\" Plot normalised MNIST images \"\"\"\n    fig = plt.figure(figsize=(10, 5))\n\n    n_tensors = len(tensors)\n    n_cols = n_tensors\n    n_rows = 1\n\n    for row in range(n_rows):\n        for col in range(n_cols):\n            idx = n_cols * row + col\n\n            if idx > n_tensors - 1:\n                break\n\n            ax = fig.add_subplot(n_rows, n_cols, idx + 1)\n            tensor = tensors[idx]\n\n            # Clip image values\n            tensor[tensor < 0] = 0\n            tensor[tensor > 1] = 1\n\n            # Put the channel last\n            tensor = torch.permute(tensor, (1, 2, 0))\n\n            ax.imshow(tensor)\n\n    plt.tight_layout()\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:34.548371Z","iopub.execute_input":"2022-09-13T21:34:34.549213Z","iopub.status.idle":"2022-09-13T21:34:34.565514Z","shell.execute_reply.started":"2022-09-13T21:34:34.549170Z","shell.execute_reply":"2022-09-13T21:34:34.564010Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_data(data_loader):\n    \"\"\" Plot data used to supposedly train the target model \"\"\"\n    images = []\n\n    for _, (data, _) in enumerate(data_loader):\n        for i in range(5):\n            images.append(data[i])\n\n        break\n\n    plot_images(images)\n\nplot_data(attacker_loader_train)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:34.568926Z","iopub.execute_input":"2022-09-13T21:34:34.570808Z","iopub.status.idle":"2022-09-13T21:34:35.246173Z","shell.execute_reply.started":"2022-09-13T21:34:34.570771Z","shell.execute_reply":"2022-09-13T21:34:35.245243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Target Model\n\nBuild a PyTorch Model that queries the API as its prediction step.","metadata":{}},{"cell_type":"code","source":"## Target model\nclass TargetModel(torch.nn.Module):\n    def forward(self, x):\n\n        # Put the channel last\n        x = torch.permute(x, (0, 2, 3, 1))\n\n        # Query the Target Model to get a confidence score\n        try:\n            response = requests.post(\"https://inference.fly.dev/score\", data = json.dumps({\"input\": x.tolist()}))\n            res = json.loads(response.text)\n        except requests.exceptions.ConnectionError:\n            print(\"Connection problems. Contact the CTF organizers.\")\n\n        return np.array(res['outputs'])\n\n# Get target\ntarget_model = TargetModel()","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:35.247747Z","iopub.execute_input":"2022-09-13T21:34:35.248120Z","iopub.status.idle":"2022-09-13T21:34:35.255158Z","shell.execute_reply.started":"2022-09-13T21:34:35.248083Z","shell.execute_reply":"2022-09-13T21:34:35.253946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_target_labels(loader):\n    \"\"\" Gather data labels from the target model \"\"\"\n    attacker_targets = []\n\n    for data, label in tqdm(loader):\n        data.float()\n\n        with torch.no_grad():\n            attacker_targets.append(target_model(data))\n\n    return torch.tensor(np.concatenate(attacker_targets))\n\nattacker_targets_train = get_target_labels(attacker_loader_train)\nattacker_targets_test = get_target_labels(attacker_loader_test)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:35:41.669983Z","iopub.execute_input":"2022-09-13T21:35:41.671638Z","iopub.status.idle":"2022-09-13T21:35:47.640043Z","shell.execute_reply.started":"2022-09-13T21:35:41.671590Z","shell.execute_reply":"2022-09-13T21:35:47.639095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Saved it for you to reuse it ;)\nnp.savez('./labels_train.npz', attacker_targets_train.numpy())\nnp.savez('./labels_test.npz', attacker_targets_test.numpy())","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:36:02.605744Z","iopub.execute_input":"2022-09-13T21:36:02.606165Z","iopub.status.idle":"2022-09-13T21:36:02.614027Z","shell.execute_reply.started":"2022-09-13T21:36:02.606130Z","shell.execute_reply":"2022-09-13T21:36:02.613036Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def reconstruct_samples(model):\n    \"\"\" Visualize the model's capacity to reconstruct the original image based on the Target model output \"\"\"\n    images = []\n\n    for batch, (data, _) in enumerate(attacker_loader_train):\n        # Take the model output intercepted from the target model\n        # For train data\n        targets = attacker_targets_train[\n            batch*attacker_hyperparams.batch_size:\n            (batch+1)*attacker_hyperparams.batch_size\n            ].float()\n\n        for i in range(2):\n\n            with torch.no_grad():\n                reconstructed_image = model(torch.unsqueeze(targets[i], 0).cuda()).squeeze(0)\n\n            images.append(data[i])\n            images.append(reconstructed_image.cpu())\n        break\n\n    for batch, (data, _) in enumerate(attacker_loader_test):\n        # Take the model output intercepted from the target model\n        # For test data\n        targets = attacker_targets_test[\n            batch*attacker_hyperparams.batch_size:\n            (batch+1)*attacker_hyperparams.batch_size\n            ].float()\n\n        for i in range(2):\n\n            with torch.no_grad():\n                reconstructed_image = model(torch.unsqueeze(targets[i], 0).cuda()).squeeze(0)\n\n            images.append(data[i])\n            images.append(reconstructed_image.cpu())\n        break\n\n    plot_images(images)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:36:05.332496Z","iopub.execute_input":"2022-09-13T21:36:05.333112Z","iopub.status.idle":"2022-09-13T21:36:05.342498Z","shell.execute_reply.started":"2022-09-13T21:36:05.333076Z","shell.execute_reply":"2022-09-13T21:36:05.341266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Attack Model\n\nAn inversion attack is easiest when the ratio of input size (in other words the number outputs of the target model) to output size (in other words the size of the original image) is high. In our case we upscale a tensor of size 6 into an image of size (32x32).","metadata":{}},{"cell_type":"code","source":"## Attack Model\nclass AttackModel(torch.nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.layers = torch.nn.Sequential(\n            torch.nn.Linear(6, 16*16*4),\n            torch.nn.ReLU(),\n            torch.nn.Unflatten(1, (4, 16, 16)),\n            torch.nn.ConvTranspose2d(\n                in_channels=4,\n                out_channels=1,\n                kernel_size=3,\n                padding=1,\n                stride=2,\n                output_padding=1),\n        )\n\n    def forward(self, x):\n        return self.layers(x)\n\n# Get attacker\nattacker = AttackModel().cuda()\nsummary(attacker, (6, ))","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:36:08.392750Z","iopub.execute_input":"2022-09-13T21:36:08.393638Z","iopub.status.idle":"2022-09-13T21:36:16.868304Z","shell.execute_reply.started":"2022-09-13T21:36:08.393600Z","shell.execute_reply":"2022-09-13T21:36:16.867267Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def attack(attack_model):\n    \"\"\" For the 6 output classes of the model, get its reconstructed train data sample \"\"\"\n    images = []\n\n    for idx in range(6):\n        target_output = np.zeros(6)\n        target_output[idx] = 1\n\n        with torch.no_grad():\n            reconstructed_image = attack_model(torch.tensor([target_output]).float().cuda()).squeeze(0)\n\n        images.append(reconstructed_image.cpu())\n\n    plot_images(images)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:36:20.938421Z","iopub.execute_input":"2022-09-13T21:36:20.939431Z","iopub.status.idle":"2022-09-13T21:36:20.946113Z","shell.execute_reply.started":"2022-09-13T21:36:20.939386Z","shell.execute_reply":"2022-09-13T21:36:20.945166Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inversion Attack\n\nNow onto the training of the Attack Model. The most critical part is the Loss function. **I minimize the Euclidian distance between the original image (fed to the Target Model) and the reconstructed image.**\n\nI did not use any data augmentation because the assumption that the attack model should be robust to a small transform to the original image (fed to the target model) is quite strong. Since we don't know much about the dataset used, I figured it was dangerous but do let me know if that is wise.","metadata":{}},{"cell_type":"code","source":"torch.manual_seed(0)\n\n# Train on the entire dataset\nattacker = AttackModel().cuda()\n\nhyperparams = namedtuple('hyperparams', 'batch_size,epochs,learning_rate')\n\nattacker_hyperparams = hyperparams(\n    batch_size=128,\n    epochs=100,\n    learning_rate=5e-2,\n)\n\n# Dataloader\nattacker_loader_train = torch.utils.data.DataLoader(attacker_dataset_train, batch_size=attacker_hyperparams.batch_size)\n\n# Optimizer\nattack_optim = torch.optim.Adam(attacker.parameters(), lr=attacker_hyperparams.learning_rate)\n\n# Learning Rate Schedule\nscheduler = StepLR(attack_optim, step_size=30, gamma=0.1, verbose=False)\n\nfor epoch in trange(attacker_hyperparams.epochs):\n    epoch_loss = 0.0\n\n    for batch, (data, targets) in enumerate(attacker_loader_train):\n        if not (epoch % 20) and not batch:\n            ## Generated images\n            reconstruct_samples(attacker)\n            attack(attacker)\n\n        data.float()\n        # Take the model output intercepted from the target model\n        targets = attacker_targets_train[\n            batch*attacker_hyperparams.batch_size:\n            (batch+1)*attacker_hyperparams.batch_size\n            ].float()\n\n        attack_optim.zero_grad()\n\n        # Reconstruct image\n        output = attacker(targets.cuda())\n\n        # Compute loss & Backprop\n        loss = ((output - data.cuda())**2).mean() # We want our reconstructed image to look as much like the original image as possible\n        loss.backward()\n        attack_optim.step()\n\n        # Logging/Display training progression\n        ## Loss\n        epoch_loss += output.shape[0] * loss.item()\n\n    print(epoch, epoch_loss / len(attacker_dataset_train))\n    scheduler.step()\n\nprint('Finished Training')\nattack(attacker)","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:36:23.770330Z","iopub.execute_input":"2022-09-13T21:36:23.770892Z","iopub.status.idle":"2022-09-13T21:36:27.520417Z","shell.execute_reply.started":"2022-09-13T21:36:23.770853Z","shell.execute_reply":"2022-09-13T21:36:27.519475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"try:\n    r = requests.post(\"http://flags.fly.dev:5000/score\", json=json.dumps({\"challenge_id\": \"inference\", \"submission\": \"FLAG{D3FC0N}\"}))\n    print(r.text)\nexcept requests.exceptions.ConnectionError:\n    print(\"Connection problems. Contact the CTF organizers.\")","metadata":{"execution":{"iopub.status.busy":"2022-09-13T21:34:35.447193Z","iopub.status.idle":"2022-09-13T21:34:35.448152Z","shell.execute_reply.started":"2022-09-13T21:34:35.447905Z","shell.execute_reply":"2022-09-13T21:34:35.447930Z"},"trusted":true},"execution_count":null,"outputs":[]}]}