├── .gitignore ├── requirements.txt ├── README.md ├── data.py ├── model.py ├── train_classifier.py └── train_inversion.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | .DS_Store 3 | *.pyc 4 | venv/ 5 | test.sh 6 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | cycler==0.10.0 2 | decorator==4.4.0 3 | imageio==2.5.0 4 | kiwisolver==1.1.0 5 | matplotlib==3.0.3 6 | networkx==2.3 7 | numpy==1.17.0 8 | Pillow==6.1.0 9 | pkg-resources==0.0.0 10 | pyparsing==2.4.2 11 | python-dateutil==2.8.0 12 | PyWavelets==1.0.3 13 | scikit-image==0.15.0 14 | scipy==1.3.1 15 | six==1.12.0 16 | torch==1.2.0 17 | torchvision==0.4.0 18 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Adversarial Model Inversion Attack 2 | 3 | This repo provides an example of the adversarial model inversion attack in the 4 | paper ["Neural Network Inversion in Adversarial Setting via Background Knowledge Alignment"](https://dl.acm.org/citation.cfm?id=3354261) 5 | 6 | ## Data 7 | 8 | The target classifier (identity classification) is trained on the [FaceScrub](http://vintage.winklerbros.net/facescrub.html) 9 | dataset, and the adversary will use the [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) dataset as background 10 | knowledge to train the inversion model. 11 | 12 | #### Download 13 | 14 | [FaceScrub](http://vintage.winklerbros.net/facescrub.html), [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) 15 | 16 | #### Extract Face 17 | 18 | [FaceScrub](http://vintage.winklerbros.net/facescrub.html): Extract the face of each image according to the official 19 | bounding box information. 20 | 21 | [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html): To extract the face of each image, crop the official 22 | align-cropped version (size 178 × 218) by width and height of 108 from upper left coordinate (35, 70). Please contact 23 | the authors of CelebA for details about the face identities, and then "clean" the CelebA by removing celebrities that 24 | are included in FaceScrub. 25 | 26 | Transform both datasets to greyscale images with each pixel value in [0, 1]. Resize both datasets to 64 × 64. 27 | 28 | ## Setup 29 | 30 | The code is written in Python3. You can install the required packages by running: 31 | 32 | ``` 33 | $ pip3 install -r requirements.txt 34 | ``` 35 | 36 | ## Run 37 | 38 | Train the target classifier: 39 | 40 | ``` 41 | $ python3 train_classifier.py 42 | ``` 43 | 44 | Train the inversion model: 45 | 46 | ``` 47 | $ python3 train_inversion.py 48 | ``` 49 | 50 | You can set the truncation size by the `--truncation` parameter. 51 | 52 | ## Citation 53 | 54 | ``` 55 | @inproceedings{Yang:2019:NNI:3319535.3354261, 56 | author = {Yang, Ziqi and Zhang, Jiyi and Chang, Ee-Chien and Liang, Zhenkai}, 57 | title = {Neural Network Inversion in Adversarial Setting via Background Knowledge Alignment}, 58 | booktitle = {Proceedings of the 2019 ACM SIGSAC Conference on Computer and Communications Security}, 59 | series = {CCS '19}, 60 | year = {2019}, 61 | isbn = {978-1-4503-6747-9}, 62 | location = {London, United Kingdom}, 63 | pages = {225--240}, 64 | numpages = {16}, 65 | url = {http://doi.acm.org/10.1145/3319535.3354261}, 66 | doi = {10.1145/3319535.3354261}, 67 | acmid = {3354261}, 68 | publisher = {ACM}, 69 | address = {New York, NY, USA}, 70 | keywords = {deep learning, model inversion, neural networks, privacy, security}, 71 | } 72 | ``` 73 | -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, division 2 | import os 3 | import numpy as np 4 | from torch.utils.data import Dataset 5 | from PIL import Image 6 | 7 | class FaceScrub(Dataset): 8 | def __init__(self, root, transform=None, target_transform=None, train=True): 9 | self.root = os.path.expanduser(root) 10 | self.transform = transform 11 | self.target_transform = target_transform 12 | 13 | input = np.load(os.path.join(self.root, 'facescrub.npz')) 14 | actor_images = input['actor_images'] 15 | actor_labels = input['actor_labels'] 16 | actress_images = input['actress_images'] 17 | actress_labels = input['actress_labels'] 18 | 19 | data = np.concatenate([actor_images, actress_images], axis=0) 20 | labels = np.concatenate([actor_labels, actress_labels], axis=0) 21 | 22 | v_min = data.min(axis=0) 23 | v_max = data.max(axis=0) 24 | data = (data - v_min) / (v_max - v_min) 25 | 26 | np.random.seed(666) 27 | perm = np.arange(len(data)) 28 | np.random.shuffle(perm) 29 | data = data[perm] 30 | labels = labels[perm] 31 | 32 | if train: 33 | self.data = data[0:int(0.8 * len(data))] 34 | self.labels = labels[0:int(0.8 * len(data))] 35 | else: 36 | self.data = data[int(0.8 * len(data)):] 37 | self.labels = labels[int(0.8 * len(data)):] 38 | 39 | def __len__(self): 40 | return len(self.data) 41 | 42 | def __getitem__(self, index): 43 | img, target = self.data[index], self.labels[index] 44 | img = Image.fromarray(img) 45 | 46 | if self.transform is not None: 47 | img = self.transform(img) 48 | 49 | if self.target_transform is not None: 50 | target = self.target_transform(target) 51 | 52 | return img, target 53 | 54 | class CelebA(Dataset): 55 | def __init__(self, root, transform=None, target_transform=None): 56 | self.root = os.path.expanduser(root) 57 | self.transform = transform 58 | self.target_transform = target_transform 59 | 60 | data = [] 61 | for i in range(10): 62 | data.append(np.load(os.path.join(self.root, 'celebA_64_{}.npy').format(i + 1))) 63 | data = np.concatenate(data, axis=0) 64 | 65 | v_min = data.min(axis=0) 66 | v_max = data.max(axis=0) 67 | data = (data - v_min) / (v_max - v_min) 68 | labels = np.array([0] * len(data)) 69 | 70 | self.data = data 71 | self.labels = labels 72 | 73 | def __len__(self): 74 | return len(self.data) 75 | 76 | def __getitem__(self, index): 77 | img, target = self.data[index], self.labels[index] 78 | img = Image.fromarray(img) 79 | 80 | if self.transform is not None: 81 | img = self.transform(img) 82 | 83 | if self.target_transform is not None: 84 | target = self.target_transform(target) 85 | 86 | return img, target 87 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | class Classifier(nn.Module): 7 | def __init__(self, nc, ndf, nz): 8 | super(Classifier, self).__init__() 9 | 10 | self.nc = nc 11 | self.ndf = ndf 12 | self.nz = nz 13 | 14 | self.encoder = nn.Sequential( 15 | # input is (nc) x 64 x 64 16 | nn.Conv2d(nc, ndf, 3, 1, 1), 17 | nn.BatchNorm2d(ndf), 18 | nn.MaxPool2d(2, 2, 0), 19 | nn.ReLU(True), 20 | # state size. (ndf) x 32 x 32 21 | nn.Conv2d(ndf, ndf * 2, 3, 1, 1), 22 | nn.BatchNorm2d(ndf * 2), 23 | nn.MaxPool2d(2, 2, 0), 24 | nn.ReLU(True), 25 | # state size. (ndf*2) x 16 x 16 26 | nn.Conv2d(ndf * 2, ndf * 4, 3, 1, 1), 27 | nn.BatchNorm2d(ndf * 4), 28 | nn.MaxPool2d(2, 2, 0), 29 | nn.ReLU(True), 30 | # state size. (ndf*4) x 8 x 8 31 | nn.Conv2d(ndf * 4, ndf * 8, 3, 1, 1), 32 | nn.BatchNorm2d(ndf * 8), 33 | nn.MaxPool2d(2, 2, 0), 34 | nn.ReLU(True), 35 | # state size. (ndf*8) x 4 x 4 36 | ) 37 | 38 | self.fc = nn.Sequential( 39 | nn.Linear(ndf * 8 * 4 * 4, nz * 5), 40 | nn.Dropout(0.5), 41 | nn.Linear(nz * 5, nz), 42 | ) 43 | 44 | def forward(self, x, release=False): 45 | 46 | x = x.view(-1, 1, 64, 64) 47 | x = self.encoder(x) 48 | x = x.view(-1, self.ndf * 8 * 4 * 4) 49 | x = self.fc(x) 50 | 51 | if release: 52 | return F.softmax(x, dim=1) 53 | else: 54 | return F.log_softmax(x, dim=1) 55 | 56 | class Inversion(nn.Module): 57 | def __init__(self, nc, ngf, nz, truncation, c): 58 | super(Inversion, self).__init__() 59 | 60 | self.nc = nc 61 | self.ngf = ngf 62 | self.nz = nz 63 | self.truncation = truncation 64 | self.c = c 65 | 66 | self.decoder = nn.Sequential( 67 | # input is Z 68 | nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0), 69 | nn.BatchNorm2d(ngf * 8), 70 | nn.Tanh(), 71 | # state size. (ngf*8) x 4 x 4 72 | nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1), 73 | nn.BatchNorm2d(ngf * 4), 74 | nn.Tanh(), 75 | # state size. (ngf*4) x 8 x 8 76 | nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1), 77 | nn.BatchNorm2d(ngf * 2), 78 | nn.Tanh(), 79 | # state size. (ngf*2) x 16 x 16 80 | nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1), 81 | nn.BatchNorm2d(ngf), 82 | nn.Tanh(), 83 | # state size. (ngf) x 32 x 32 84 | nn.ConvTranspose2d(ngf, nc, 4, 2, 1), 85 | nn.Sigmoid() 86 | # state size. (nc) x 64 x 64 87 | ) 88 | 89 | def forward(self, x): 90 | topk, indices = torch.topk(x, self.truncation) 91 | topk = torch.clamp(torch.log(topk), min=-1000) + self.c 92 | topk_min = topk.min(1, keepdim=True)[0] 93 | topk = topk + F.relu(-topk_min) 94 | x = torch.zeros(len(x), self.nz).cuda().scatter_(1, indices, topk) 95 | 96 | x = x.view(-1, self.nz, 1, 1) 97 | x = self.decoder(x) 98 | x = x.view(-1, 1, 64, 64) 99 | return x -------------------------------------------------------------------------------- /train_classifier.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import argparse 3 | import torch 4 | import torch.nn as nn 5 | import torch.optim as optim 6 | from torchvision import transforms 7 | import os 8 | import torch.nn.functional as F 9 | from data import FaceScrub 10 | from model import Classifier 11 | 12 | # Training settings 13 | parser = argparse.ArgumentParser(description='Adversarial Model Inversion Demo') 14 | parser.add_argument('--batch-size', type=int, default=128, metavar='') 15 | parser.add_argument('--test-batch-size', type=int, default=1000, metavar='') 16 | parser.add_argument('--epochs', type=int, default=100, metavar='') 17 | parser.add_argument('--lr', type=float, default=0.01, metavar='') 18 | parser.add_argument('--momentum', type=float, default=0.5, metavar='') 19 | parser.add_argument('--no-cuda', action='store_true', default=False) 20 | parser.add_argument('--seed', type=int, default=1, metavar='') 21 | parser.add_argument('--log-interval', type=int, default=10, metavar='') 22 | parser.add_argument('--nc', type=int, default=1) 23 | parser.add_argument('--ndf', type=int, default=128) 24 | parser.add_argument('--nz', type=int, default=530) 25 | parser.add_argument('--num_workers', type=int, default=1, metavar='') 26 | 27 | def train(classifier, log_interval, device, data_loader, optimizer, epoch): 28 | classifier.train() 29 | for batch_idx, (data, target) in enumerate(data_loader): 30 | data, target = data.to(device), target.to(device) 31 | optimizer.zero_grad() 32 | output = classifier(data) 33 | loss = F.nll_loss(output, target) 34 | loss.backward() 35 | optimizer.step() 36 | 37 | if batch_idx % log_interval == 0: 38 | print('Train Epoch: {} [{}/{}]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), 39 | len(data_loader.dataset), loss.item())) 40 | 41 | def test(classifier, device, data_loader): 42 | classifier.eval() 43 | test_loss = 0 44 | correct = 0 45 | with torch.no_grad(): 46 | for data, target in data_loader: 47 | data, target = data.to(device), target.to(device) 48 | output = classifier(data) 49 | test_loss += F.nll_loss(output, target, reduction='sum').item() 50 | pred = output.max(1, keepdim=True)[1] 51 | correct += pred.eq(target.view_as(pred)).sum().item() 52 | 53 | test_loss /= len(data_loader.dataset) 54 | print('\nTest classifier: Average loss: {:.6f}, Accuracy: {}/{} ({:.4f}%)\n'.format( 55 | test_loss, correct, len(data_loader.dataset), 100. * correct / len(data_loader.dataset))) 56 | return correct / len(data_loader.dataset) 57 | 58 | def main(): 59 | args = parser.parse_args() 60 | print("================================") 61 | print(args) 62 | print("================================") 63 | os.makedirs('out', exist_ok=True) 64 | 65 | use_cuda = not args.no_cuda and torch.cuda.is_available() 66 | device = torch.device("cuda" if use_cuda else "cpu") 67 | kwargs = {'num_workers': args.num_workers, 'pin_memory': True} if use_cuda else {} 68 | 69 | torch.manual_seed(args.seed) 70 | 71 | transform = transforms.Compose([transforms.ToTensor()]) 72 | train_set = FaceScrub('./data/facescrub', transform=transform, train=True) 73 | test_set = FaceScrub('./data/facescrub', transform=transform, train=False) 74 | 75 | train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs) 76 | test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.test_batch_size, shuffle=False, **kwargs) 77 | 78 | classifier = nn.DataParallel(Classifier(nc=args.nc, ndf=args.ndf, nz=args.nz)).to(device) 79 | optimizer = optim.Adam(classifier.parameters(), lr=0.0002, betas=(0.5, 0.999), amsgrad=True) 80 | 81 | best_cl_acc = 0 82 | best_cl_epoch = 0 83 | 84 | # Train classifier 85 | for epoch in range(1, args.epochs + 1): 86 | train(classifier, args.log_interval, device, train_loader, optimizer, epoch) 87 | cl_acc = test(classifier, device, test_loader) 88 | 89 | if cl_acc > best_cl_acc: 90 | best_cl_acc = cl_acc 91 | best_cl_epoch = epoch 92 | state = { 93 | 'epoch': epoch, 94 | 'model': classifier.state_dict(), 95 | 'optimizer': optimizer.state_dict(), 96 | 'best_cl_acc': best_cl_acc, 97 | } 98 | torch.save(state, 'out/classifier.pth') 99 | 100 | print("Best classifier: epoch {}, acc {:.4f}".format(best_cl_epoch, best_cl_acc)) 101 | 102 | if __name__ == '__main__': 103 | main() 104 | -------------------------------------------------------------------------------- /train_inversion.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import argparse 3 | import torch 4 | import torch.nn as nn 5 | import torch.optim as optim 6 | from torchvision import transforms 7 | import os, shutil 8 | from data import FaceScrub, CelebA 9 | from model import Classifier, Inversion 10 | import torch.nn.functional as F 11 | import torchvision.utils as vutils 12 | 13 | # Training settings 14 | parser = argparse.ArgumentParser(description='Adversarial Model Inversion Demo') 15 | parser.add_argument('--batch-size', type=int, default=128, metavar='') 16 | parser.add_argument('--test-batch-size', type=int, default=1000, metavar='') 17 | parser.add_argument('--epochs', type=int, default=100, metavar='') 18 | parser.add_argument('--lr', type=float, default=0.01, metavar='') 19 | parser.add_argument('--momentum', type=float, default=0.5, metavar='') 20 | parser.add_argument('--no-cuda', action='store_true', default=False) 21 | parser.add_argument('--seed', type=int, default=1, metavar='') 22 | parser.add_argument('--log-interval', type=int, default=10, metavar='') 23 | parser.add_argument('--nc', type=int, default=1) 24 | parser.add_argument('--ndf', type=int, default=128) 25 | parser.add_argument('--ngf', type=int, default=128) 26 | parser.add_argument('--nz', type=int, default=530) 27 | parser.add_argument('--truncation', type=int, default=530) 28 | parser.add_argument('--c', type=float, default=50.) 29 | parser.add_argument('--num_workers', type=int, default=1, metavar='') 30 | 31 | def train(classifier, inversion, log_interval, device, data_loader, optimizer, epoch): 32 | classifier.eval() 33 | inversion.train() 34 | 35 | for batch_idx, (data, target) in enumerate(data_loader): 36 | data, target = data.to(device), target.to(device) 37 | optimizer.zero_grad() 38 | with torch.no_grad(): 39 | prediction = classifier(data, release=True) 40 | reconstruction = inversion(prediction) 41 | loss = F.mse_loss(reconstruction, data) 42 | loss.backward() 43 | optimizer.step() 44 | 45 | if batch_idx % log_interval == 0: 46 | print('Train Epoch: {} [{}/{}]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), 47 | len(data_loader.dataset), loss.item())) 48 | 49 | def test(classifier, inversion, device, data_loader, epoch, msg): 50 | classifier.eval() 51 | inversion.eval() 52 | mse_loss = 0 53 | plot = True 54 | with torch.no_grad(): 55 | for data, target in data_loader: 56 | data, target = data.to(device), target.to(device) 57 | prediction = classifier(data, release=True) 58 | reconstruction = inversion(prediction) 59 | mse_loss += F.mse_loss(reconstruction, data, reduction='sum').item() 60 | 61 | if plot: 62 | truth = data[0:32] 63 | inverse = reconstruction[0:32] 64 | out = torch.cat((inverse, truth)) 65 | for i in range(4): 66 | out[i * 16:i * 16 + 8] = inverse[i * 8:i * 8 + 8] 67 | out[i * 16 + 8:i * 16 + 16] = truth[i * 8:i * 8 + 8] 68 | vutils.save_image(out, 'out/recon_{}_{}.png'.format(msg.replace(" ", ""), epoch), normalize=False) 69 | plot = False 70 | 71 | mse_loss /= len(data_loader.dataset) * 64 * 64 72 | print('\nTest inversion model on {} set: Average MSE loss: {:.6f}\n'.format(msg, mse_loss)) 73 | return mse_loss 74 | 75 | def main(): 76 | args = parser.parse_args() 77 | print("================================") 78 | print(args) 79 | print("================================") 80 | os.makedirs('out', exist_ok=True) 81 | 82 | use_cuda = not args.no_cuda and torch.cuda.is_available() 83 | device = torch.device("cuda" if use_cuda else "cpu") 84 | kwargs = {'num_workers': args.num_workers, 'pin_memory': True} if use_cuda else {} 85 | 86 | torch.manual_seed(args.seed) 87 | 88 | transform = transforms.Compose([transforms.ToTensor()]) 89 | train_set = CelebA('./data/celebA', transform=transform) 90 | # Inversion attack on TRAIN data of facescrub classifier 91 | test1_set = FaceScrub('./data/facescrub', transform=transform, train=True) 92 | # Inversion attack on TEST data of facescrub classifier 93 | test2_set = FaceScrub('./data/facescrub', transform=transform, train=False) 94 | 95 | train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, **kwargs) 96 | test1_loader = torch.utils.data.DataLoader(test1_set, batch_size=args.test_batch_size, shuffle=False, **kwargs) 97 | test2_loader = torch.utils.data.DataLoader(test2_set, batch_size=args.test_batch_size, shuffle=False, **kwargs) 98 | 99 | classifier = nn.DataParallel(Classifier(nc=args.nc, ndf=args.ndf, nz=args.nz)).to(device) 100 | inversion = nn.DataParallel(Inversion(nc=args.nc, ngf=args.ngf, nz=args.nz, truncation=args.truncation, c=args.c)).to(device) 101 | optimizer = optim.Adam(inversion.parameters(), lr=0.0002, betas=(0.5, 0.999), amsgrad=True) 102 | 103 | # Load classifier 104 | path = 'out/classifier.pth' 105 | try: 106 | checkpoint = torch.load(path) 107 | classifier.load_state_dict(checkpoint['model']) 108 | epoch = checkpoint['epoch'] 109 | best_cl_acc = checkpoint['best_cl_acc'] 110 | print("=> loaded classifier checkpoint '{}' (epoch {}, acc {:.4f})".format(path, epoch, best_cl_acc)) 111 | except: 112 | print("=> load classifier checkpoint '{}' failed".format(path)) 113 | return 114 | 115 | # Train inversion model 116 | best_recon_loss = 999 117 | for epoch in range(1, args.epochs + 1): 118 | train(classifier, inversion, args.log_interval, device, train_loader, optimizer, epoch) 119 | recon_loss = test(classifier, inversion, device, test1_loader, epoch, 'test1') 120 | test(classifier, inversion, device, test2_loader, epoch, 'test2') 121 | 122 | if recon_loss < best_recon_loss: 123 | best_recon_loss = recon_loss 124 | state = { 125 | 'epoch': epoch, 126 | 'model': inversion.state_dict(), 127 | 'optimizer': optimizer.state_dict(), 128 | 'best_recon_loss': best_recon_loss 129 | } 130 | torch.save(state, 'out/inversion.pth') 131 | shutil.copyfile('out/recon_test1_{}.png'.format(epoch), 'out/best_test1.png') 132 | shutil.copyfile('out/recon_test2_{}.png'.format(epoch), 'out/best_test2.png') 133 | 134 | if __name__ == '__main__': 135 | main() 136 | --------------------------------------------------------------------------------