Model_Inversion_Attack.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Model Inversion Attack"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "694cdd24278509a0"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Settings and Imports"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9d3437ecb2094261"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# suppress warnings\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "#autoreload other packages when code changed\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "37ba81d5975f3f37"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from tqdm.notebook import tqdm"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f81e680cbdce67e1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "#Own Code\n",
    "from privacyflow.configs import path_configs\n",
    "from privacyflow.datasets import faces_dataset, mi_dataset\n",
    "from privacyflow.models import face_models, cifar_models, cifar_autoencoder, celeba_autoencoder"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ce76f169ed9fd0e1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "#Check if GPU is available\n",
    "if torch.cuda.is_available():\n",
    "    print(\"GPU will be used\")\n",
    "    device = torch.device('cuda')\n",
    "else:\n",
    "    print(\"No GPU available\")\n",
    "    device = torch.device('cpu')"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7008b94693789a9a"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## CIFAR-10 - Denoising Autoencoder"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fcaacb93a2478251"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# #Custom Transformation for adding noise to trtaining_data\n",
    "# class AddGaussianNoise(object):\n",
    "#     def __init__(self, mean:float=0.0, std:float=0.0):\n",
    "#         self.std = std\n",
    "#         self.mean = mean\n",
    "# \n",
    "#     def __call__(self, tensor):\n",
    "#         tensor = tensor + torch.randn(tensor.size()) * self.std + self.mean\n",
    "#         return torch.clip(tensor,min=0.0,max=1.0)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "14389853d57c092e"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "autoencoder_transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    # AddGaussianNoise(mean=0.0, std=0.1)\n",
    "])\n",
    "cifar10_ds_train_autoencoder = torchvision.datasets.CIFAR10(root=path_configs.CIFAR_FOLDER_PATH,\n",
    "                                                            transform=autoencoder_transform,\n",
    "                                                            train=True,\n",
    "                                                            download=True)\n",
    "\n",
    "cifar10_ds_test_autoencoder = torchvision.datasets.CIFAR10(root=path_configs.CIFAR_FOLDER_PATH,\n",
    "                                                           transform=autoencoder_transform,\n",
    "                                                           train=False,\n",
    "                                                           download=True)\n",
    "\n",
    "#Combine the datasets for the usage for the autoencoder\n",
    "cifar10_ds_autoencoder = torch.utils.data.ConcatDataset([cifar10_ds_test_autoencoder, cifar10_ds_train_autoencoder])\n",
    "cifar10_dl_autoencoder = DataLoader(cifar10_ds_autoencoder, batch_size=64, shuffle=True)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "feb5e1e7342bc6a8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def train_autoencoder_no_logs(model: nn.Module,\n",
    "                        train_dl: torch.utils.data.DataLoader,\n",
    "                        optimizer: torch.optim,\n",
    "                        criterion: nn.Module,\n",
    "                        num_epochs: int = 15):\n",
    "    model.train()\n",
    "    model = model.to(device)\n",
    "    for _ in tqdm(range(num_epochs),leave=False):\n",
    "        for images, _ in train_dl:\n",
    "            images = images.to(device)\n",
    "            model_inputs = torch.clip(images + torch.rand(images.size(),device=device) * 0.3, min=0.0, max=1.0)\n",
    "            optimizer.zero_grad()\n",
    "            model_outputs = model(model_inputs)\n",
    "            loss = criterion(model_outputs, images)\n",
    "            loss.backward()\n",
    "            optimizer.step()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9b33bc950a44b808"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "autoencoder = cifar_autoencoder.CifarDenoisingAutoencoder()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(autoencoder.parameters(), lr=0.0001)\n",
    "\n",
    "train_autoencoder_no_logs(autoencoder,\n",
    "                    train_dl=cifar10_ds_autoencoder,\n",
    "                    optimizer=optimizer,\n",
    "                    criterion=criterion,\n",
    "                    num_epochs=20)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3951c3aa0aab2b3d"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## CIFAR-10 Modell Inversion Attacke"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "458eb003c331dad6"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def reconstruction_attack_cifar10(\n",
    "        model:nn.Module,\n",
    "        autoencoder:nn.Module,\n",
    "        start_tensor:torch.Tensor,\n",
    "        target:torch.Tensor,\n",
    "        num_epochs:int=10_000,\n",
    "        learning_rate:float =0.01,\n",
    "        use_autoencoder:bool=True) -> torch.Tensor:\n",
    "    \n",
    "    #Params for reconstruction attack\n",
    "    criterion_tensor = nn.NLLLoss()\n",
    "    for _ in tqdm(range(num_epochs),leave=False):\n",
    "        optimizer_tensor = torch.optim.Adam([start_tensor],lr=learning_rate)\n",
    "        #Update tensor due to model\n",
    "        optimizer_tensor.zero_grad()\n",
    "        #the exp functions is due to the fact, that the cifar10 models are using log_softmax as last activation function\n",
    "        # output = torch.exp(model(start_tensor))\n",
    "        output = model(start_tensor)\n",
    "        loss = criterion_tensor(output,target)\n",
    "        loss.backward()\n",
    "        optimizer_tensor.step()\n",
    "        if use_autoencoder:\n",
    "            start_tensor = torch.tensor(autoencoder(start_tensor),device=device,requires_grad=True)\n",
    "    return start_tensor"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e0f833b2769f80d4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "attacked_model = cifar_models.CifarCNNModel()\n",
    "attacked_model.load_state_dict(torch.load(f\"{path_configs.MODELS_TRAINED_BASE_PATH}/cifar_10_base.pl\"))\n",
    "\n",
    "#set start tensor with random values with size equal to image size\n",
    "#start_tensor = torch.rand([3,32,32],device=device).unsqueeze(0).requires_grad_()\n",
    "start_tensor = torch.empty(3,32,32,device=device).fill_(0.5).unsqueeze(0).requires_grad_()\n",
    "target_label = torch.tensor([8],dtype=torch.long,device=device)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7dab3549cf5651"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "autoencoder = autoencoder.to(device)\n",
    "attacked_model = attacked_model.to(device)\n",
    "\n",
    "\n",
    "torchvision.transforms.ToPILImage()(start_tensor.squeeze(0)).show()\n",
    "\n",
    "reconstructed_image = reconstruction_attack_cifar10(model=attacked_model,\n",
    "                                            autoencoder=autoencoder,\n",
    "                                            start_tensor=start_tensor,\n",
    "                                            target=target_label,\n",
    "                                            num_epochs=10_000,\n",
    "                                            learning_rate=0.01,\n",
    "                                            use_autoencoder=False)\n",
    "torchvision.transforms.ToPILImage()(reconstructed_image.squeeze(0)).show()\n",
    "\n",
    "\n",
    "reconstructed_image = reconstruction_attack_cifar10(model=attacked_model,\n",
    "                                            autoencoder=autoencoder,\n",
    "                                            start_tensor=reconstructed_image,\n",
    "                                            target=target_label,\n",
    "                                            num_epochs=10_000,\n",
    "                                            learning_rate=0.1,\n",
    "                                            use_autoencoder=True)\n",
    "torchvision.transforms.ToPILImage()(reconstructed_image.squeeze(0)).show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "16f2796956859b81"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## CelebA Denoising Autoencoder"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "db4386d53df15716"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "celeba_autoencoder_transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize((224, 224)),\n",
    "    torchvision.transforms.ToTensor()\n",
    "])\n",
    "\n",
    "train_dataset_celeba = faces_dataset.FacesDataset(label_cols='all', \n",
    "                                               mode=\"train\",\n",
    "                                               transform=celeba_autoencoder_transform)\n",
    "val_dataset_celeba = faces_dataset.FacesDataset(label_cols='all', \n",
    "                                             mode=\"val\", \n",
    "                                             transform=celeba_autoencoder_transform)\n",
    "test_dataset_celeba = faces_dataset.FacesDataset(label_cols='all', \n",
    "                                              mode=\"test\", \n",
    "                                              transform=celeba_autoencoder_transform)\n",
    "\n",
    "#Combien Datasets for training of autoencoder\n",
    "dataset_celeba_combines = torch.utils.data.ConcatDataset([train_dataset_celeba,val_dataset_celeba,test_dataset_celeba])\n",
    "dl_celeba_combines = DataLoader(dataset_celeba_combines,batch_size=32,num_workers=4,shuffle=True)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7f2d15fdf477d2d2"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "autoencoder = celeba_autoencoder.CelebADenoisingAutoencoder()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(autoencoder.parameters(), lr=0.0001)\n",
    "\n",
    "train_autoencoder_no_logs(autoencoder,\n",
    "                    train_dl=dl_celeba_combines,\n",
    "                    optimizer=optimizer,\n",
    "                    criterion=criterion,\n",
    "                    num_epochs=10)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "98ef5d31f78ac8cb"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "img = dataset_celeba_combines[128][0]\n",
    "torchvision.transforms.ToPILImage()(img).save(\"./privacyflow/images/faces_autoencoder1.jpg\")\n",
    "\n",
    "img = torch.clip(img + torch.rand(img.size()) * 0.3, min=0.0, max=1.0)\n",
    "torchvision.transforms.ToPILImage()(img).save(\"./privacyflow/images/faces_autoencoder2.jpg\")\n",
    "\n",
    "autoencoder = autoencoder.to('cpu')\n",
    "img = autoencoder(img)\n",
    "torchvision.transforms.ToPILImage()(img).save(\"./privacyflow/images/faces_autoencoder3.jpg\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1002b2f7f7d0884f"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def reconstruction_attack_cifar10(\n",
    "        model:nn.Module,\n",
    "        autoencoder:nn.Module,\n",
    "        start_tensor:torch.Tensor,\n",
    "        target:torch.Tensor,\n",
    "        num_epochs:int=10_000,\n",
    "        learning_rate:float =0.01,\n",
    "        use_autoencoder:bool=True) -> torch.Tensor:\n",
    "    \n",
    "    #Params for reconstruction attack\n",
    "    criterion_tensor = nn.BCELoss()\n",
    "    for _ in tqdm(range(num_epochs),leave=False):\n",
    "        optimizer_tensor = torch.optim.Adam([start_tensor],lr=learning_rate)\n",
    "        #Update tensor due to model\n",
    "        optimizer_tensor.zero_grad()\n",
    "        output = model(start_tensor)\n",
    "        loss = criterion_tensor(output,target)\n",
    "        loss.backward()\n",
    "        optimizer_tensor.step()\n",
    "        if use_autoencoder:\n",
    "            start_tensor = torch.tensor(autoencoder(start_tensor),device=device,requires_grad=True)\n",
    "    return start_tensor"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "10339eacc36a8eda"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "attacked_model = face_models.get_FaceModelResNet(40)\n",
    "attacked_model.load_state_dict(torch.load(f\"{path_configs.MODELS_TRAINED_BASE_PATH}/face_base_model.pl\"))\n",
    "\n",
    "#set start tensor with random values with size equal to image size\n",
    "#start_tensor = torch.rand([3,224,224],device=device).unsqueeze(0).requires_grad_()\n",
    "start_tensor = torch.empty(3,224,224,device=device).fill_(0.5).unsqueeze(0).requires_grad_()\n",
    "target_label = dataset_celeba_combines[128][1].unsqueeze(0).to(device)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "756e958d5289b4b3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "autoencoder = autoencoder.to(device)\n",
    "attacked_model = attacked_model.to(device)\n",
    "\n",
    "torchvision.transforms.ToPILImage()(start_tensor.squeeze(0)).show()\n",
    "\n",
    "reconstructed_image = reconstruction_attack_cifar10(model=attacked_model,\n",
    "                                            autoencoder=autoencoder,\n",
    "                                            start_tensor=start_tensor,\n",
    "                                            target=target_label,\n",
    "                                            num_epochs=10_000,\n",
    "                                            learning_rate=0.01,\n",
    "                                            use_autoencoder=False)\n",
    "torchvision.transforms.ToPILImage()(reconstructed_image.squeeze(0)).show()\n",
    "\n",
    "\n",
    "reconstructed_image = reconstruction_attack_cifar10(model=attacked_model,\n",
    "                                            autoencoder=autoencoder,\n",
    "                                            start_tensor=reconstructed_image,\n",
    "                                            target=target_label,\n",
    "                                            num_epochs=10_000,\n",
    "                                            learning_rate=0.1,\n",
    "                                            use_autoencoder=True)\n",
    "torchvision.transforms.ToPILImage()(reconstructed_image.squeeze(0)).show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4847e12cfaf21f2d"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "ffa8d1c9c1c1d528"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}