├── .gitignore
├── Attacking and Protecting Data Privacy in Edge-Cloud Collaborative Inference Systems.pdf
├── Model Inversion Attacks Against Collaborative Inference.pdf
├── README.md
├── defenses.sh
├── inverse_blackbox_decoder_CIFAR.py
├── inverse_query_free_CIFAR.py
├── inverse_whitebox_CIFAR.py
├── inverse_whitebox_CIFAR_defense.py
├── inverse_whitebox_MNIST.py
├── inverse_whitebox_MNIST_defense.py
├── net.py
├── net.pyc
├── noise_generation.py
├── noise_generation_opt.py
├── training.py
├── utils.py
└── utils.pyc


/.gitignore:
--------------------------------------------------------------------------------
 1 | .DS_Store
 2 | *.pyc
 3 | *.npy
 4 | data/
 5 | checkpoints/
 6 | inverted_whitebox/
 7 | inverted_access_free/
 8 | inverted_blackbox_decoder/
 9 | noise/
10 | 


--------------------------------------------------------------------------------
/Attacking and Protecting Data Privacy in Edge-Cloud Collaborative Inference Systems.pdf:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zechenghe/Inverse_Collaborative_Inference/f010aed0035238c573037b04c47c334ea733d15c/Attacking and Protecting Data Privacy in Edge-Cloud Collaborative Inference Systems.pdf


--------------------------------------------------------------------------------
/Model Inversion Attacks Against Collaborative Inference.pdf:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zechenghe/Inverse_Collaborative_Inference/f010aed0035238c573037b04c47c334ea733d15c/Model Inversion Attacks Against Collaborative Inference.pdf


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Model Inversion Attack against Collaborative Inference
 2 | 
 3 | This code implements model inversion attacks against collaborative inference in the following two papers:
 4 | 
 5 | Zecheng He, Tianwei Zhang and Ruby Lee, "Model Inversion Attacks Against Collaborative Inference", 35th Annual Computer Security Applications Conference (ACSAC'19), San Juan, Dec 2019 ([paper1](https://github.com/zechenghe/Inverse_Collaborative_Inference/blob/master/Model%20Inversion%20Attacks%20Against%20Collaborative%20Inference.pdf))
 6 | 
 7 | 
 8 | Zecheng He, Tianwei Zhang and Ruby Lee, "Attacking and Protecting Data Privacy in Edge-Cloud Collaborative Inference Systems", IEEE Internet of Things Journal (IoTJ), 2020 ([paper2](https://github.com/zechenghe/Inverse_Collaborative_Inference/blob/master/Attacking%20and%20Protecting%20Data%20Privacy%20in%20Edge-Cloud%20Collaborative%20Inference%20Systems.pdf))
 9 | 
10 | In [paper1](https://github.com/zechenghe/Inverse_Collaborative_Inference/blob/master/Model%20Inversion%20Attacks%20Against%20Collaborative%20Inference.pdf), we provide three attacks, i.e. rMSE (Section 4), blackbox inverse network (Section 5) and query-free attack (Section 6) on CIFAR10 dataset. Attacks against MNIST are similar.
11 | 
12 | In [paper2](https://github.com/zechenghe/Inverse_Collaborative_Inference/blob/master/Attacking%20and%20Protecting%20Data%20Privacy%20in%20Edge-Cloud%20Collaborative%20Inference%20Systems.pdf), we provide adding noise and dropout as defenses (Section IV).
13 | 
14 | ### 1.Dependencies:
15 | #### python 2.7
16 | #### numpy
17 |     pip install numpy
18 | #### pytorch 1.0.0
19 |     pip install torch
20 | #### torchvision version 0.2.1:
21 |     pip install torchvision==0.2.1
22 | 
23 | <br/>
24 | 
25 | ### 2.Run the code:
26 | #### (1) Train the target CIFAR model to inverse
27 | 
28 |     python training.py --dataset CIFAR10 --epochs 50
29 | 
30 | #### (2) Whitebox Regularized Maximum Likelihood Estimation (rMLE, Section 4)
31 | 
32 |     python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer ReLU22 --lambda_TV 1e1 --lambda_l2 0.0
33 | 
34 | #### (3) Blackbox Inverse Network (Section 5)
35 | #### Train inverse network
36 |     python inverse_blackbox_decoder_CIFAR.py --training --layer ReLU22 --iter 50 --decodername CIFAR10CNNDecoderReLU22
37 | #### Inference inverse network
38 |     python inverse_blackbox_decoder_CIFAR.py --testing --decodername CIFAR10CNNDecoderReLU22 --layer ReLU22
39 | 
40 | #### (4) Query-free Attack (Section 6)
41 | 
42 | #### Train a shadow model
43 |     python inverse_query_free_CIFAR.py --training --layer ReLU22 --iter 50
44 | 
45 | #### Inverse the shadow model
46 |     python inverse_query_free_CIFAR.py --testing --layer ReLU22 --iter 500 --learning_rate 1e-1 --lambda_TV 2e0 --lambda_l2 0.0
47 | 
48 | #### (5) Defense (Paper 2 Section IV)
49 | 
50 |     ./defense.sh
51 | 
52 | <br/>
53 | 
54 | ### 3.Tips
55 | 
56 | (1) Please make sure to use torchvision v0.2.1:
57 | 
58 |     import torchvision
59 |     print torchvision.__version__
60 | 
61 | (2) If no gpu supported on your machine, add --nogpu option in the command line.
62 | 
63 | (3) Please feel free to add --novalidation in your command line if it takes you too long to run in cpu-only mode and the model training/inverse are conducted on the same machine. It will disable evalTest() and evalTestSplitModel(), which are only used for validating the pre-trained models.
64 | 
65 | <br/>
66 | 
67 | ### 4.Reference
68 | You are encouraged to cite the following papers.
69 | ```
70 | @inproceedings{he2019model,
71 |   title={Model inversion attacks against collaborative inference},
72 |   author={He, Zecheng and Zhang, Tianwei and Lee, Ruby B},
73 |   booktitle={Proceedings of the 35th Annual Computer Security Applications Conference},
74 |   pages={148--162},
75 |   year={2019},
76 |   organization={ACM}
77 | }
78 | 
79 | @article{he2020attacking,
80 |   title={Attacking and Protecting Data Privacy in Edge-Cloud Collaborative Inference Systems},
81 |   author={He, Zecheng and Zhang, Tianwei and Lee, Ruby B},
82 |   journal={IEEE Internet of Things Journal},
83 |   year={2020},
84 |   publisher={IEEE}
85 | }
86 | ```
87 | 


--------------------------------------------------------------------------------
/defenses.sh:
--------------------------------------------------------------------------------
 1 | #! /bin/bash
 2 | 
 3 | # Random noise
 4 | Random_Noise () {
 5 |     python inverse_whitebox_MNIST_defense.py --noise_type Gaussian --layer ReLU2 --add_noise_to_input
 6 |     python inverse_whitebox_MNIST_defense.py --noise_type Gaussian --layer ReLU2
 7 |     python inverse_whitebox_MNIST_defense.py --noise_type Laplace --layer ReLU2 --add_noise_to_input
 8 |     python inverse_whitebox_MNIST_defense.py --noise_type Laplace --layer ReLU2
 9 | }
10 | 
11 | # Dropout
12 | Dropout (){
13 |     python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer ReLU2 --add_noise_to_input
14 |     python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer ReLU2
15 | }
16 | 
17 | Dropout_CIFAR (){
18 |     python2 inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer ReLU22 --iters 5000 --learning_rate 1e-2 --lambda_TV 1e1 --lambda_l2 0.0
19 |     python2 inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer ReLU22 --add_noise_to_input --iters 5000 --learning_rate 1e-2 --lambda_TV 1e1 --lambda_l2 0.0
20 |     python2 inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer ReLU12 --iters 5000 --learning_rate 1e-2 --lambda_TV 0.0 --lambda_l2 0.0
21 |     python2 inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer pool1 --iters 5000 --learning_rate 1e-2 --lambda_TV 0.0 --lambda_l2 0.0
22 |     python2 inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer conv22 --iters 5000 --learning_rate 1e-2 --lambda_TV 1e1 --lambda_l2 0.0
23 |     python2 inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer pool2 --iters 5000 --learning_rate 1e-2 --lambda_TV 1e1 --lambda_l2 0.0
24 | }
25 | 
26 | # Opt noise generation
27 | Opt_Noise_Generation (){
28 |     python noise_generation_opt.py --noise_sourceLayer pool1 --noise_targetLayer conv2
29 |     python inverse_whitebox_MNIST_defense.py --noise_type noise_gen_opt --layer pool1 --noise_targetLayer conv2
30 | 
31 |     python noise_generation_opt.py --noise_sourceLayer conv2 --noise_targetLayer ReLU2
32 |     python inverse_whitebox_MNIST_defense.py --noise_type noise_gen_opt --layer conv2 --noise_targetLayer ReLU2
33 | 
34 |     python noise_generation_opt.py --noise_sourceLayer ReLU2 --noise_targetLayer pool2
35 |     python inverse_whitebox_MNIST_defense.py --noise_type noise_gen_opt --layer ReLU2 --noise_targetLayer pool2
36 | 
37 |     python noise_generation_opt.py --noise_sourceLayer pool2 --noise_targetLayer fc1
38 |     python inverse_whitebox_MNIST_defense.py --noise_type noise_gen_opt --layer pool2 --noise_targetLayer fc1
39 | }
40 | 
41 | #Random_Noise
42 | #Dropout
43 | #Opt_Noise_Generation
44 | 
45 | #python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer pool1
46 | #python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer conv2
47 | #python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer ReLU2
48 | #python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer pool2
49 | #python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer fc1
50 | 
51 | Dropout_CIFAR
52 | 


--------------------------------------------------------------------------------
/inverse_blackbox_decoder_CIFAR.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2019-09-01
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | 
 17 | from net import *
 18 | from utils import *
 19 | 
 20 | #####################
 21 | # Training:
 22 | # python inverse_blackbox_decoder_CIFAR.py --layer ReLU22 --iter 50 --training --decodername CIFAR10CNNDecoderReLU22
 23 | #
 24 | # Testing:
 25 | # python inverse_blackbox_decoder_CIFAR.py --testing --decodername CIFAR10CNNDecoderReLU22 --layer ReLU22
 26 | #####################
 27 | 
 28 | 
 29 | def trainDecoderDNN(DATASET = 'CIFAR10', network = 'CIFAR10CNNDecoder', NEpochs = 50, imageWidth = 32,
 30 |             imageHeight = 32, imageSize = 32*32, NChannels = 3, NClasses = 10, layer = 'ReLU22', BatchSize = 32, learningRate = 1e-3,
 31 |             NDecreaseLR = 20, eps = 1e-3, AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth", save_decoder_dir = "checkpoints/CIFAR10/",
 32 |             decodername_name = 'CIFAR10CNNDecoderReLU22', gpu = True, validation=False):
 33 | 
 34 |     print "DATASET: ", DATASET
 35 | 
 36 |     if DATASET == 'CIFAR10':
 37 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
 38 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
 39 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 40 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 41 |         tsf = {
 42 |             'train': transforms.Compose(
 43 |             [
 44 |             transforms.ToTensor(),
 45 |             Normalize
 46 |             ]),
 47 |             'test': transforms.Compose(
 48 |             [
 49 |             transforms.ToTensor(),
 50 |             Normalize
 51 |             ])
 52 |         }
 53 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
 54 |                                         download=True, transform = tsf['train'])
 55 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
 56 |                                        download=True, transform = tsf['test'])
 57 | 
 58 |     else:
 59 |         print "Dataset unsupported"
 60 |         exit(1)
 61 | 
 62 |     print "len(trainset) ", len(trainset)
 63 |     print "len(testset) ", len(testset)
 64 |     x_train, y_train = trainset.data, trainset.targets,
 65 |     x_test, y_test = testset.data, testset.targets,
 66 | 
 67 |     print "x_train.shape ", x_train.shape
 68 |     print "x_test.shape ", x_test.shape
 69 | 
 70 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = BatchSize,
 71 |                                       shuffle = False, num_workers = 1)
 72 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 73 |                                       shuffle = False, num_workers = 1)
 74 |     trainIter = iter(trainloader)
 75 |     testIter = iter(testloader)
 76 | 
 77 |     net = torch.load(model_dir + model_name)
 78 |     net.eval()
 79 |     print "Validate the model accuracy..."
 80 | 
 81 |     if validation:
 82 |         accTest = evalTest(testloader, net, gpu = gpu)
 83 | 
 84 |     # Get dims of input/output, and construct the network
 85 |     batchX, batchY = trainIter.next()
 86 |     if gpu:
 87 |         batchX = batchX.cuda()
 88 |     originalModelOutput = net.getLayerOutput(batchX, net.layerDict[layer]).clone()
 89 | 
 90 |     decoderNetDict = {
 91 |         'CIFAR10CNNDecoder':{
 92 |             'conv11': CIFAR10CNNDecoderconv1,
 93 |             'ReLU22': CIFAR10CNNDecoderReLU22,
 94 |             'ReLU32': CIFAR10CNNDecoderReLU32
 95 |         }
 96 |     }
 97 |     decoderNetFunc = decoderNetDict[network][layer]
 98 |     decoderNet = decoderNetFunc(originalModelOutput.shape[1])
 99 |     if gpu:
100 |         decoderNet = decoderNet.cuda()
101 | 
102 |     print decoderNet
103 | 
104 |     NBatch = len(trainset) / BatchSize
105 |     MSELossLayer = torch.nn.MSELoss()
106 |     if gpu:
107 |         MSELossLayer = MSELossLayer.cuda()
108 | 
109 |     # Find the optimal config according to the hardware
110 |     cudnn.benchmark = True
111 |     optimizer = optim.Adam(params = decoderNet.parameters(), lr = learningRate, eps = eps, amsgrad = AMSGrad)
112 | 
113 |     for epoch in range(NEpochs):
114 |         lossTrain = 0.0
115 |         accTrain = 0.0
116 |         for i in range(NBatch):
117 |             try:
118 |                 batchX, batchY = trainIter.next()
119 |             except StopIteration:
120 |                 trainIter = iter(trainloader)
121 |                 batchX, batchY = trainIter.next()
122 | 
123 |             if gpu:
124 |                 batchX = batchX.cuda()
125 |                 batchY = batchY.cuda()
126 | 
127 |             optimizer.zero_grad()
128 | 
129 |             originalModelOutput = net.getLayerOutput(batchX, net.layerDict[layer]).clone()
130 |             decoderNetOutput = decoderNet.forward(originalModelOutput)
131 | 
132 |             assert batchX.cpu().detach().numpy().shape == decoderNetOutput.cpu().detach().numpy().shape
133 | 
134 |             featureLoss = MSELossLayer(batchX, decoderNetOutput)
135 |             totalLoss = featureLoss
136 |             totalLoss.backward()
137 |             optimizer.step()
138 | 
139 |             lossTrain += totalLoss / NBatch
140 | 
141 |         valData, valLabel = iter(testloader).next()
142 |         if gpu:
143 |             valData = valData.cuda()
144 |             valLabel = valLabel.cuda()
145 |         originalModelOutput = net.getLayerOutput(valData, net.layerDict[layer]).clone()
146 |         decoderNetOutput = decoderNet.forward(originalModelOutput)
147 |         valLoss = MSELossLayer(valData, decoderNetOutput)
148 | 
149 |         print "Epoch ", epoch, "Train Loss: ", lossTrain.cpu().detach().numpy(), "Test Loss: ", valLoss.cpu().detach().numpy()
150 |         if (epoch + 1) % NDecreaseLR == 0:
151 |             learningRate = learningRate / 2.0
152 |             setLearningRate(optimizer, learningRate)
153 | 
154 |     if validation:
155 |         accTestEnd = evalTest(testloader, net, gpu = gpu)
156 |         if accTest != accTestEnd:
157 |             print "Something wrong. Original model has been modified!"
158 |             exit(1)
159 | 
160 |     if not os.path.exists(save_decoder_dir):
161 |         os.makedirs(save_decoder_dir)
162 |     torch.save(decoderNet, save_decoder_dir + decodername_name)
163 |     print "Model saved"
164 | 
165 |     newNet = torch.load(save_decoder_dir + decodername_name)
166 |     newNet.eval()
167 |     print "Model restore done"
168 | 
169 | 
170 | def inverse(DATASET = 'CIFAR10', imageWidth = 32, inverseClass = None, imageHeight = 32,
171 |         imageSize = 32*32, NChannels = 3, NClasses = 10, layer = 'conv11',
172 |         model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth", decoder_name = "CIFAR10CNNDecoderconv11.pth",
173 |         save_img_dir = "inverted_blackbox_decoder/CIFAR10/", gpu = True, validation=False):
174 | 
175 |     print "DATASET: ", DATASET
176 |     print "inverseClass: ", inverseClass
177 | 
178 |     assert inverseClass < NClasses
179 | 
180 |     if DATASET == 'CIFAR10':
181 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
182 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
183 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
184 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
185 |         tsf = {
186 |             'train': transforms.Compose(
187 |             [
188 |             transforms.ToTensor(),
189 |             Normalize
190 |             ]),
191 |             'test': transforms.Compose(
192 |             [
193 |             transforms.ToTensor(),
194 |             Normalize
195 |             ])
196 |         }
197 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
198 |                                         download=True, transform = tsf['train'])
199 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
200 |                                        download=True, transform = tsf['test'])
201 | 
202 |     else:
203 |         print "Dataset unsupported"
204 |         exit(1)
205 | 
206 | 
207 |     print "len(trainset) ", len(trainset)
208 |     print "len(testset) ", len(testset)
209 |     x_train, y_train = trainset.train_data, trainset.train_labels,
210 |     x_test, y_test = testset.test_data, testset.test_labels,
211 | 
212 |     print "x_train.shape ", x_train.shape
213 |     print "x_test.shape ", x_test.shape
214 | 
215 | 
216 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
217 |                                       shuffle = False, num_workers = 1)
218 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
219 |                                       shuffle = False, num_workers = 1)
220 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
221 |                                       shuffle = False, num_workers = 1)
222 |     trainIter = iter(trainloader)
223 |     testIter = iter(testloader)
224 |     inverseIter = iter(inverseloader)
225 | 
226 | 
227 |     net = torch.load(model_dir + model_name)
228 |     if not gpu:
229 |         net = net.cpu()
230 |     net.eval()
231 |     print "Validate the model accuracy..."
232 | 
233 |     if validation:
234 |         accTest = evalTest(testloader, net, gpu = gpu)
235 | 
236 |     decoderNet = torch.load(model_dir + decoder_name)
237 |     if not gpu:
238 |         decoderNet = decoderNet.cpu()
239 |     decoderNet.eval()
240 |     print decoderNet
241 |     print "Validate the alternative model..."
242 |     batchX, batchY = iter(testloader).next()
243 |     if gpu:
244 |         batchX = batchX.cuda()
245 |         batchY = batchY.cuda()
246 | 
247 |     print "batchX.shape ", batchX.cpu().detach().numpy().shape
248 | 
249 |     MSELossLayer = torch.nn.MSELoss()
250 |     if gpu:
251 |         MSELossLayer = MSELossLayer.cuda()
252 |     originalModelOutput = net.getLayerOutput(batchX, net.layerDict[layer]).clone()
253 |     decoderNetOutput = decoderNet.forward(originalModelOutput)
254 | 
255 |     assert batchX.cpu().detach().numpy().shape == decoderNetOutput.cpu().detach().numpy().shape
256 |     print "decoderNetOutput.shape ", decoderNetOutput.cpu().detach().numpy().shape
257 |     print "MSE ", MSELossLayer(batchX, decoderNetOutput).cpu().detach().numpy()
258 | 
259 |     targetImg, _ = getImgByClass(inverseIter, C = inverseClass)
260 |     print "targetImg.size()", targetImg.size()
261 | 
262 |     deprocessImg = deprocess(targetImg.clone())
263 | 
264 |     if not os.path.exists(save_img_dir):
265 |         os.makedirs(save_img_dir)
266 |     torchvision.utils.save_image(deprocessImg, save_img_dir + str(inverseClass) + '-ref.png')
267 | 
268 |     if gpu:
269 |         targetImg = targetImg.cuda()
270 |     targetLayer = net.layerDict[layer]
271 |     refFeature = net.getLayerOutput(targetImg, targetLayer)
272 | 
273 |     print "refFeature.size()", refFeature.size()
274 | 
275 |     xGen = decoderNet.forward(refFeature)
276 |     print "MSE ", MSELossLayer(targetImg, xGen).cpu().detach().numpy()
277 | 
278 |     # save the final result
279 |     imgGen = xGen.clone()
280 |     imgGen = deprocess(imgGen)
281 |     torchvision.utils.save_image(imgGen, save_img_dir + str(inverseClass) + '-inv.png')
282 | 
283 |     print "Done"
284 | 
285 | 
286 | if __name__ == '__main__':
287 |     import argparse
288 |     import sys
289 |     import traceback
290 | 
291 |     try:
292 |         parser = argparse.ArgumentParser()
293 |         parser.add_argument('--dataset', type = str, default = 'CIFAR10')
294 |         parser.add_argument('--network', type = str, default = 'CIFAR10CNNDecoder')
295 | 
296 |         parser.add_argument('--training', dest='training', action='store_true')
297 |         parser.add_argument('--testing', dest='training', action='store_false')
298 |         parser.set_defaults(training=False)
299 | 
300 |         parser.add_argument('--iters', type = int, default = 500)
301 |         parser.add_argument('--eps', type = float, default = 1e-3)
302 |         parser.add_argument('--AMSGrad', type = bool, default = True)
303 |         parser.add_argument('--batch_size', type = int, default = 32)
304 |         parser.add_argument('--learning_rate', type = float, default = 1e-3)
305 |         parser.add_argument('--decrease_LR', type = int, default = 10)
306 |         parser.add_argument('--layer', type = str, default = 'ReLU22')
307 |         parser.add_argument('--save_iter', type = int, default = 10)
308 |         parser.add_argument('--inverseClass', type = int, default = 0)
309 |         parser.add_argument('--decodername', type = str, default = "CIFAR10CNNDecoderReLU22")
310 | 
311 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
312 |         parser.set_defaults(gpu=True)
313 | 
314 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
315 |         parser.set_defaults(validation=True)
316 |         args = parser.parse_args()
317 | 
318 |         model_dir = "checkpoints/" + args.dataset + '/'
319 |         model_name = "ckpt.pth"
320 |         decoder_name = args.decodername + '.pth'
321 | 
322 |         save_img_dir = "inverted_blackbox_decoder/" + args.dataset + '/' + args.layer + '/'
323 | 
324 |         if args.dataset == 'CIFAR10':
325 | 
326 |             imageWidth = 32
327 |             imageHeight = 32
328 |             imageSize = imageWidth * imageHeight
329 |             NChannels = 3
330 |             NClasses = 10
331 | 
332 |         else:
333 |             print "No Dataset Found"
334 |             exit()
335 | 
336 |         if args.training:
337 |             trainDecoderDNN(DATASET = args.dataset, network = 'CIFAR10CNNDecoder', NEpochs = args.iters, imageWidth = imageWidth,
338 |             imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer, BatchSize = args.batch_size, learningRate = args.learning_rate,
339 |             NDecreaseLR = args.decrease_LR, eps = args.eps, AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth", save_decoder_dir = "checkpoints/CIFAR10/",
340 |             decodername_name = decoder_name, gpu = args.gpu, validation=args.validation)
341 | 
342 |         else:
343 |             for c in range(NClasses):
344 |                 inverse(DATASET = args.dataset, imageHeight = imageHeight, imageWidth = imageWidth, inverseClass = c,
345 |                 imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer,
346 |                 model_dir = model_dir, model_name = model_name, decoder_name = decoder_name,
347 |                 save_img_dir = save_img_dir, gpu = args.gpu, validation=args.validation)
348 | 
349 |     except:
350 |         traceback.print_exc(file=sys.stdout)
351 |         sys.exit(1)
352 | 


--------------------------------------------------------------------------------
/inverse_query_free_CIFAR.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2019-08-31
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | 
 17 | from net import *
 18 | from utils import *
 19 | 
 20 | #####################
 21 | # Training:
 22 | # python inverse_access_free_CIFAR.py --layer ReLU22 --iter 50 --training
 23 | #
 24 | # Testing:
 25 | # python inverse_access_free_CIFAR.py --testing --layer ReLU22 --iter 500 --learning_rate 1e-1 --lambda_TV 2e0 --lambda_l2 0.0
 26 | #####################
 27 | 
 28 | 
 29 | def trainAlternativeDNN(DATASET = 'CIFAR10', network = 'CIFAR10CNNAlternative', NEpochs = 200, imageWidth = 32,
 30 |         imageHeight = 32, imageSize = 32*32, NChannels = 3, NClasses = 10, layer = 'conv', BatchSize = 32, learningRate = 1e-3,
 31 |         NDecreaseLR = 5, eps = 1e-3, AMSGrad = True, model_dir = "checkpoints/MNIST/", model_name = "ckpt.pth", save_alternative_model_dir = "checkpoints/MNIST/",
 32 |         alternative_model_name = "LeNetAccessFree.pth", gpu = True, validation=False):
 33 | 
 34 |     print "DATASET: ", DATASET
 35 | 
 36 |     if DATASET == 'CIFAR10':
 37 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
 38 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
 39 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 40 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 41 |         tsf = {
 42 |             'train': transforms.Compose(
 43 |             [
 44 |             transforms.ToTensor(),
 45 |             Normalize
 46 |             ]),
 47 |             'test': transforms.Compose(
 48 |             [
 49 |             transforms.ToTensor(),
 50 |             Normalize
 51 |             ])
 52 |         }
 53 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
 54 |                                         download=True, transform = tsf['train'])
 55 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
 56 |                                        download=True, transform = tsf['test'])
 57 | 
 58 |     else:
 59 |         print "Dataset unsupported"
 60 |         exit(1)
 61 | 
 62 |     print "len(trainset) ", len(trainset)
 63 |     print "len(testset) ", len(testset)
 64 |     x_train, y_train = trainset.data, trainset.targets,
 65 |     x_test, y_test = testset.data, testset.targets,
 66 | 
 67 |     print "x_train.shape ", x_train.shape
 68 |     print "x_test.shape ", x_test.shape
 69 | 
 70 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = BatchSize,
 71 |                                       shuffle = False, num_workers = 1)
 72 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 73 |                                       shuffle = False, num_workers = 1)
 74 |     trainIter = iter(trainloader)
 75 |     testIter = iter(testloader)
 76 | 
 77 |     # Load the trained model
 78 |     net = torch.load(model_dir + model_name)
 79 |     if not gpu:
 80 |         net = net.cpu()
 81 | 
 82 |     net.eval()
 83 |     print "Validate the model accuracy..."
 84 | 
 85 |     if validation:
 86 |         accTest = evalTest(testloader, net, gpu = gpu)
 87 | 
 88 |     altnetDict = {
 89 |         'CIFAR10CNNAlternative':{
 90 |             'conv11': CIFAR10CNNAlternativeconv11,
 91 |             'ReLU22': CIFAR10CNNAlternativeReLU22,
 92 |             'ReLU32': CIFAR10CNNAlternativeReLU32
 93 |         }
 94 |     }
 95 |     alternativeNetFunc = altnetDict[network][layer]
 96 | 
 97 |     if gpu:
 98 |         alternativeNet = alternativeNetFunc(NChannels).cuda()
 99 |     else:
100 |         alternativeNet = alternativeNetFunc(NChannels)
101 | 
102 |     print alternativeNet
103 | 
104 |     # Get dims of input/output, and construct the network
105 |     batchX, batchY = trainIter.next()
106 |     if gpu:
107 |         batchX = batchX.cuda()
108 | 
109 |     edgeOutput = alternativeNet.forward(batchX)
110 | 
111 |     if gpu:
112 |         cloudOuput = net.forward_from(edgeOutput, layer).clone()
113 |     else:
114 |         cloudOuput = net.forward_from(edgeOutput, layer)
115 | 
116 |     print "edgeOutput.size", edgeOutput.size()
117 |     print "cloudOuput.size", cloudOuput.size()
118 | 
119 |     NBatch = len(trainset) / BatchSize
120 |     if gpu:
121 |         CrossEntropyLossLayer = nn.CrossEntropyLoss().cuda()
122 |     else:
123 |         CrossEntropyLossLayer = nn.CrossEntropyLoss()
124 | 
125 |     # Find the optimal config according to the hardware
126 |     cudnn.benchmark = True
127 |     optimizer = optim.Adam(params = alternativeNet.parameters(), lr = learningRate, eps = eps, amsgrad = AMSGrad)
128 | 
129 |     # Sanity check
130 |     valData, valLabel = iter(testloader).next()
131 |     if gpu:
132 |         valData = valData.cuda()
133 |         valLabel = valLabel.cuda()
134 | 
135 |     edgeOutput = alternativeNet.forward(valData)
136 |     cloudOuput = net.forward_from(edgeOutput, layer).clone()
137 |     valLoss = CrossEntropyLossLayer(cloudOuput, valLabel)
138 | 
139 |     print "Test Loss without training: ", valLoss.cpu().detach().numpy()
140 | 
141 |     for epoch in range(NEpochs):
142 |         lossTrain = 0.0
143 |         accTrain = 0.0
144 |         for i in range(NBatch):
145 |             try:
146 |                 batchX, batchY = trainIter.next()
147 |             except StopIteration:
148 |                 trainIter = iter(trainloader)
149 |                 batchX, batchY = trainIter.next()
150 | 
151 |             if gpu:
152 |                 batchX = batchX.cuda()
153 |                 batchY = batchY.cuda()
154 | 
155 |             optimizer.zero_grad()
156 | 
157 |             edgeOutput = alternativeNet.forward(batchX).clone()
158 |             cloudOuput = net.forward_from(edgeOutput, layer)
159 | 
160 |             featureLoss = CrossEntropyLossLayer(cloudOuput, batchY)
161 | 
162 |             totalLoss = featureLoss
163 |             totalLoss.backward()
164 |             optimizer.step()
165 | 
166 |             lossTrain += totalLoss.cpu().detach().numpy() / NBatch
167 | 
168 |         valData, valLabel = iter(testloader).next()
169 |         if gpu:
170 |             valData = valData.cuda()
171 |             valLabel = valLabel.cuda()
172 |         edgeOutput = alternativeNet.forward(valData)
173 |         cloudOuput = net.forward_from(edgeOutput, layer).clone()
174 |         valLoss = CrossEntropyLossLayer(cloudOuput, valLabel)
175 | 
176 |         if validation:
177 |             accTestSplitModel = evalTestSplitModel(testloader, alternativeNet, net, layer, gpu = gpu)
178 |             print "Epoch ", epoch, "Train Loss: ", lossTrain, "Test Loss: ", valLoss.cpu().detach().numpy(), "Test Accuracy", accTestSplitModel
179 | 
180 |         if (epoch + 1) % NDecreaseLR == 0:
181 |             learningRate = learningRate / 2.0
182 |             setLearningRate(optimizer, learningRate)
183 | 
184 |     if validation:
185 |         accTestEnd = evalTest(testloader, net, gpu = gpu)
186 |         if accTest != accTestEnd:
187 |             print "Something wrong. Original model has been modified!"
188 |             exit(1)
189 | 
190 |     if not os.path.exists(save_alternative_model_dir):
191 |         os.makedirs(save_alternative_model_dir)
192 |     torch.save(alternativeNet, save_alternative_model_dir + alternative_model_name)
193 |     print "Model saved"
194 | 
195 |     newNet = torch.load(save_alternative_model_dir + alternative_model_name)
196 |     newNet.eval()
197 |     print "Model restore done"
198 | 
199 | 
200 | def inverse(DATASET = 'CIFAR10', NIters = 500, imageWidth = 32, inverseClass = None,
201 |         imageHeight = 32, imageSize = 32*32, NChannels = 3, NClasses = 10, layer = 'ReLU2',
202 |         learningRate = 1e-3, NDecreaseLR = 20, eps = 1e-3, lambda_TV = 1e3, lambda_l2 = 1.0,
203 |         AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth",
204 |         alternative_model_name = "CIFAR10CNNAccessFree.pth", save_img_dir = "inverted_access_free/CIFAR10/MSE_TV/",
205 |         saveIter = 10, gpu = True, validation=False):
206 | 
207 |     print "DATASET: ", DATASET
208 |     print "inverseClass: ", inverseClass
209 | 
210 |     assert inverseClass < NClasses
211 | 
212 |     if DATASET == 'CIFAR10':
213 | 
214 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
215 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
216 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
217 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
218 |         tsf = {
219 |             'train': transforms.Compose(
220 |             [
221 |             transforms.ToTensor(),
222 |             Normalize
223 |             ]),
224 |             'test': transforms.Compose(
225 |             [
226 |             transforms.ToTensor(),
227 |             Normalize
228 |             ])
229 |         }
230 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
231 |                                         download=True, transform = tsf['train'])
232 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
233 |                                        download=True, transform = tsf['test'])
234 | 
235 |     else:
236 |         print "Dataset unsupported"
237 |         exit(1)
238 | 
239 | 
240 |     print "len(trainset) ", len(trainset)
241 |     print "len(testset) ", len(testset)
242 |     x_train, y_train = trainset.train_data, trainset.train_labels,
243 |     x_test, y_test = testset.test_data, testset.test_labels,
244 | 
245 |     print "x_train.shape ", x_train.shape
246 |     print "x_test.shape ", x_test.shape
247 | 
248 | 
249 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
250 |                                       shuffle = False, num_workers = 1)
251 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
252 |                                       shuffle = False, num_workers = 1)
253 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
254 |                                       shuffle = False, num_workers = 1)
255 |     trainIter = iter(trainloader)
256 |     testIter = iter(testloader)
257 |     inverseIter = iter(inverseloader)
258 | 
259 |     # Load the trained model
260 |     net = torch.load(model_dir + model_name)
261 |     if not gpu:
262 |         net = net.cpu()
263 |     net.eval()
264 |     print "Validate the model accuracy..."
265 | 
266 |     if validation:
267 |         accTest = evalTest(testloader, net, gpu = gpu)
268 | 
269 |     alternativeNet = torch.load(model_dir + alternative_model_name)
270 |     if not gpu:
271 |         alternativeNet = alternativeNet.cpu()
272 |     alternativeNet.eval()
273 |     print alternativeNet
274 |     #print "Validate the alternative model..."
275 |     batchX, batchY = iter(testloader).next()
276 |     if gpu:
277 |         batchX = batchX.cuda()
278 |         batchY = batchY.cuda()
279 | 
280 |     if gpu:
281 |         MSELossLayer = torch.nn.MSELoss().cuda()
282 |     else:
283 |         MSELossLayer = torch.nn.MSELoss()
284 | 
285 |     originalModelOutput = net.getLayerOutput(batchX, net.layerDict[layer]).clone()
286 |     alternativeModelOutput = alternativeNet.forward(batchX)
287 | 
288 |     print "originalModelOutput.shape: ", originalModelOutput.shape, "alternativeModelOutput.shape: ", alternativeModelOutput.shape
289 |     print "MSE difference on layer " + layer, MSELossLayer(originalModelOutput, alternativeModelOutput)
290 | 
291 | 
292 |     targetImg, _ = getImgByClass(inverseIter, C = inverseClass)
293 |     print "targetImg.size()", targetImg.size()
294 | 
295 |     deprocessImg = deprocess(targetImg.clone())
296 | 
297 |     if not os.path.exists(save_img_dir):
298 |         os.makedirs(save_img_dir)
299 |     torchvision.utils.save_image(deprocessImg, save_img_dir + str(inverseClass) + '-ref.png')
300 | 
301 |     if gpu:
302 |         targetImg = targetImg.cuda()
303 |     targetLayer = net.layerDict[layer]
304 |     refFeature = net.getLayerOutput(targetImg, targetLayer)
305 | 
306 |     print "refFeature.size()", refFeature.size()
307 | 
308 |     if gpu:
309 |         xGen = torch.zeros(targetImg.size(), requires_grad = True, device="cuda")
310 |     else:
311 |         xGen = torch.zeros(targetImg.size(), requires_grad = True)
312 | 
313 |     optimizer = optim.Adam(params = [xGen], lr = learningRate, eps = eps, amsgrad = AMSGrad)
314 | 
315 |     for i in range(NIters):
316 | 
317 |         optimizer.zero_grad()
318 | 
319 |         xFeature = alternativeNet.forward(xGen)
320 |         featureLoss = ((xFeature - refFeature)**2).mean()
321 | 
322 |         TVLoss = TV(xGen)
323 |         normLoss = l2loss(xGen)
324 | 
325 |         totalLoss = featureLoss + lambda_TV * TVLoss + lambda_l2 * normLoss
326 | 
327 |         totalLoss.backward(retain_graph=True)
328 |         optimizer.step()
329 | 
330 |         print "Iter ", i, "Feature loss: ", featureLoss.cpu().detach().numpy(), "TVLoss: ", TVLoss.cpu().detach().numpy(), "l2Loss: ", normLoss.cpu().detach().numpy()
331 | 
332 |     # save the final result
333 |     imgGen = xGen.clone()
334 |     imgGen = deprocess(imgGen)
335 |     torchvision.utils.save_image(imgGen, save_img_dir + str(inverseClass) + '-inv.png')
336 | 
337 |     print "Done"
338 | 
339 | 
340 | if __name__ == '__main__':
341 |     import argparse
342 |     import sys
343 |     import traceback
344 | 
345 |     try:
346 |         parser = argparse.ArgumentParser()
347 |         parser.add_argument('--dataset', type = str, default = 'CIFAR10')
348 |         parser.add_argument('--network', type = str, default = 'CIFAR10CNNAlternative')
349 | 
350 |         parser.add_argument('--training', dest='training', action='store_true')
351 |         parser.add_argument('--testing', dest='training', action='store_false')
352 |         parser.set_defaults(training=False)
353 |         parser.add_argument('--iters', type = int, default = 50)
354 |         parser.add_argument('--eps', type = float, default = 1e-3)
355 |         parser.add_argument('--lambda_TV', type = float, default = 1.0)
356 |         parser.add_argument('--lambda_l2', type = float, default = 1.0)
357 |         parser.add_argument('--AMSGrad', type = bool, default = True)
358 |         parser.add_argument('--batch_size', type = int, default = 32)
359 |         parser.add_argument('--learning_rate', type = float, default = 1e-3)
360 |         parser.add_argument('--decrease_LR', type = int, default = 10)
361 |         parser.add_argument('--layer', type = str, default = 'conv11')
362 |         parser.add_argument('--method', type = str, default = "MSE_TV")
363 |         parser.add_argument('--save_iter', type = int, default = 10)
364 |         parser.add_argument('--inverseClass', type = int, default = 0)
365 |         parser.add_argument('--altmodelname', type = str, default = "CIFAR10CNNAccessFree")
366 | 
367 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
368 |         parser.set_defaults(gpu=True)
369 | 
370 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
371 |         parser.set_defaults(validation=True)
372 |         args = parser.parse_args()
373 | 
374 |         model_dir = "checkpoints/" + args.dataset + '/'
375 |         model_name = "ckpt.pth"
376 |         alternative_model_name = args.altmodelname + args.layer + '.pth'
377 | 
378 |         save_img_dir = "inverted_access_free/" + args.dataset + '/' + args.layer + '/'
379 | 
380 |         if args.dataset == 'CIFAR10':
381 | 
382 |             imageWidth = 32
383 |             imageHeight = 32
384 |             imageSize = imageWidth * imageHeight
385 |             NChannels = 3
386 |             NClasses = 10
387 | 
388 |         else:
389 |             print "No Dataset Found"
390 |             exit()
391 | 
392 |         if args.training:
393 |             trainAlternativeDNN(DATASET = args.dataset, network = 'CIFAR10CNNAlternative', NEpochs = args.iters, imageWidth = imageWidth,
394 |             imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer, BatchSize = args.batch_size, learningRate = args.learning_rate,
395 |             NDecreaseLR = args.decrease_LR, eps = args.eps, AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth", save_alternative_model_dir = "checkpoints/CIFAR10/",
396 |             alternative_model_name = alternative_model_name, gpu = args.gpu, validation=args.validation)
397 | 
398 |         else:
399 |             for c in range(NClasses):
400 |                 inverse(DATASET = args.dataset, NIters = args.iters, imageWidth = imageWidth, inverseClass = c,
401 |                 imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer,
402 |                 learningRate = args.learning_rate, NDecreaseLR = args.decrease_LR, eps = args.eps, lambda_TV = args.lambda_TV, lambda_l2 = args.lambda_l2,
403 |                 AMSGrad = args.AMSGrad, model_dir = model_dir, model_name = model_name, alternative_model_name = alternative_model_name,
404 |                 save_img_dir = save_img_dir, saveIter = args.save_iter, gpu = args.gpu, validation=args.validation)
405 | 
406 |     except:
407 |         traceback.print_exc(file=sys.stdout)
408 |         sys.exit(1)
409 | 


--------------------------------------------------------------------------------
/inverse_whitebox_CIFAR.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2019-08-31
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | from net import *
 17 | from utils import *
 18 | 
 19 | #####################
 20 | # Useful Hyperparameters:
 21 | 
 22 | # CIFAR conv11
 23 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer conv11 --lambda_TV 0.0 --lambda_l2 0.0
 24 | 
 25 | # CIFAR ReLU22
 26 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer ReLU22 --lambda_TV 1e1 --lambda_l2 0.0
 27 | 
 28 | # CIFAR ReLU32
 29 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer ReLU32 --lambda_TV 1e3 --lambda_l2 0.0
 30 | 
 31 | # CIFAR fc1
 32 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-3 --layer fc1 --lambda_TV 1e4 --lambda_l2 0.0
 33 | # python inverse_whitebox_CIFAR.py --iters 1000 --learning_rate 5e-2 --layer fc1 --lambda_TV 1e5 --lambda_l2 0.0
 34 | 
 35 | # CIFAR fc2
 36 | # python inverse_whitebox_CIFAR.py --iters 500 --learning_rate 1e-1 --layer fc2 --lambda_TV 5e2 --lambda_l2 1e2
 37 | 
 38 | # CIFAR label
 39 | # python inverse_whitebox_CIFAR.py --iters 500 --learning_rate 5e-2 --layer label --lambda_TV 5e-1 --lambda_l2 0.0
 40 | #####################
 41 | 
 42 | def inverse(DATASET = 'CIFAR10', network = 'CIFAR10CNN', NIters = 500, imageWidth = 32, inverseClass = None,
 43 |         imageHeight = 32, imageSize = 32*32, NChannels = 3, NClasses = 10, layer = 'conv22',
 44 |         BatchSize = 32, learningRate = 1e-3, NDecreaseLR = 20, eps = 1e-3, lambda_TV = 1e3, lambda_l2 = 1.0,
 45 |         AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth",
 46 |         save_img_dir = "inverted/CIFAR10/MSE_TV/", saveIter = 10, gpu = True, validation=False):
 47 | 
 48 |     print "DATASET: ", DATASET
 49 |     print "inverseClass: ", inverseClass
 50 | 
 51 |     assert inverseClass < NClasses
 52 | 
 53 |     if DATASET == 'CIFAR10':
 54 | 
 55 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
 56 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
 57 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 58 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 59 | 
 60 |         tsf = {
 61 |             'train': transforms.Compose(
 62 |             [
 63 |             transforms.ToTensor(),
 64 |             Normalize
 65 |             ]),
 66 |             'test': transforms.Compose(
 67 |             [
 68 |             transforms.ToTensor(),
 69 |             Normalize
 70 |             ])
 71 |         }
 72 | 
 73 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
 74 |                                         download=True, transform = tsf['train'])
 75 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
 76 |                                        download=True, transform = tsf['test'])
 77 | 
 78 | 
 79 |     print "len(trainset) ", len(trainset)
 80 |     print "len(testset) ", len(testset)
 81 |     x_train, y_train = trainset.data, trainset.targets,
 82 |     x_test, y_test = testset.data, testset.targets,
 83 | 
 84 |     print "x_train.shape ", x_train.shape
 85 |     print "x_test.shape ", x_test.shape
 86 | 
 87 | 
 88 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
 89 |                                       shuffle = False, num_workers = 1)
 90 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 91 |                                       shuffle = False, num_workers = 1)
 92 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
 93 |                                       shuffle = False, num_workers = 1)
 94 |     trainIter = iter(trainloader)
 95 |     testIter = iter(testloader)
 96 |     inverseIter = iter(inverseloader)
 97 | 
 98 |     net = torch.load(model_dir + model_name)
 99 |     if not gpu:
100 |         net = net.cpu()
101 | 
102 |     net.eval()
103 |     print "Validate the model accuracy..."
104 |     if validation:
105 |         accTest = evalTest(testloader, net, gpu = gpu)
106 | 
107 |     targetImg, _ = getImgByClass(inverseIter, C = inverseClass)
108 |     print "targetImg.size()", targetImg.size()
109 | 
110 |     deprocessImg = deprocess(targetImg.clone())
111 | 
112 |     if not os.path.exists(save_img_dir):
113 |         os.makedirs(save_img_dir)
114 |     torchvision.utils.save_image(deprocessImg, save_img_dir + str(inverseClass) + '-ref.png')
115 | 
116 |     if gpu:
117 |         targetImg = targetImg.cuda()
118 |         softmaxLayer = nn.Softmax().cuda()
119 | 
120 |     if layer == 'prob':
121 |         reflogits = net.forward(targetImg)
122 |         refFeature = softmaxLayer(reflogits)
123 |     elif layer == 'label':
124 |         refFeature = torch.zeros(1,NClasses)
125 |         refFeature[0, inverseClass] = 1
126 |     else:
127 |         targetLayer = net.layerDict[layer]
128 |         refFeature = net.getLayerOutput(targetImg, targetLayer)
129 | 
130 |     print "refFeature.size()", refFeature.size()
131 | 
132 |     if gpu:
133 |         xGen = torch.zeros(targetImg.size(), requires_grad = True, device="cuda")
134 |     else:
135 |         xGen = torch.zeros(targetImg.size(), requires_grad = True)
136 | 
137 |     optimizer = optim.Adam(params = [xGen], lr = learningRate, eps = eps, amsgrad = AMSGrad)
138 | 
139 |     for i in range(NIters):
140 | 
141 |         optimizer.zero_grad()
142 |         if layer == 'prob':
143 |             xlogits = net.forward(xGen)
144 |             xFeature = softmaxLayer(xlogits)
145 |             featureLoss = ((xFeature - refFeature)**2).mean()
146 |         elif layer == 'label':
147 |             xlogits = net.forward(xGen)
148 |             xFeature = softmaxLayer(xlogits)
149 |             featureLoss = - torch.log(xFeature[0, inverseClass])
150 |         else:
151 |             xFeature = net.getLayerOutput(xGen, targetLayer)
152 |             featureLoss = ((xFeature - refFeature)**2).mean()
153 | 
154 |         TVLoss = TV(xGen)
155 |         normLoss = l2loss(xGen)
156 | 
157 |         totalLoss = featureLoss + lambda_TV * TVLoss + lambda_l2 * normLoss #- 1.0 * conv1Loss
158 | 
159 |         totalLoss.backward(retain_graph=True)
160 |         optimizer.step()
161 | 
162 |         print "Iter ", i, "Feature loss: ", featureLoss.cpu().detach().numpy(), "TVLoss: ", TVLoss.cpu().detach().numpy(), "l2Loss: ", normLoss.cpu().detach().numpy()
163 | 
164 |     # save the final result
165 |     imgGen = xGen.clone()
166 |     imgGen = deprocess(imgGen)
167 |     torchvision.utils.save_image(imgGen, save_img_dir + str(inverseClass) + '-inv.png')
168 | 
169 |     print "targetImg l1 Stat:"
170 |     getL1Stat(net, targetImg)
171 |     print "xGen l1 Stat:"
172 |     getL1Stat(net, xGen)
173 |     print "Done"
174 | 
175 | 
176 | if __name__ == '__main__':
177 |     import argparse
178 |     import sys
179 |     import traceback
180 | 
181 |     try:
182 |         parser = argparse.ArgumentParser()
183 |         parser.add_argument('--dataset', type = str, default = 'CIFAR10')
184 |         parser.add_argument('--network', type = str, default = 'CIFAR10CNN')
185 |         parser.add_argument('--iters', type = int, default = 500)
186 |         parser.add_argument('--eps', type = float, default = 1e-3)
187 |         parser.add_argument('--lambda_TV', type = float, default = 1.0)
188 |         parser.add_argument('--lambda_l2', type = float, default = 1.0)
189 |         parser.add_argument('--AMSGrad', type = bool, default = True)
190 |         parser.add_argument('--batch_size', type = int, default = 32)
191 |         parser.add_argument('--learning_rate', type = float, default = 1e-3)
192 |         parser.add_argument('--decrease_LR', type = int, default = 20)
193 |         parser.add_argument('--layer', type = str, default = 'conv22')
194 |         parser.add_argument('--save_iter', type = int, default = 10)
195 |         parser.add_argument('--inverseClass', type = int, default = None)
196 | 
197 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
198 |         parser.set_defaults(gpu=True)
199 | 
200 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
201 |         parser.set_defaults(validation=True)
202 |         args = parser.parse_args()
203 | 
204 |         model_dir = "checkpoints/" + args.dataset + '/'
205 |         model_name = "ckpt.pth"
206 | 
207 |         save_img_dir = "inverted_whitebox/" + args.dataset + '/' + args.layer + '/'
208 | 
209 |         if args.dataset == 'MNIST':
210 | 
211 |             imageWidth = 28
212 |             imageHeight = 28
213 |             imageSize = imageWidth * imageHeight
214 |             NChannels = 1
215 |             NClasses = 10
216 | 
217 |         elif args.dataset == 'CIFAR10':
218 | 
219 |             imageWidth = 32
220 |             imageHeight = 32
221 |             imageSize = imageWidth * imageHeight
222 |             NChannels = 3
223 |             NClasses = 10
224 | 
225 |         else:
226 |             print "No Dataset Found"
227 |             exit()
228 | 
229 |         for c in range(NClasses):
230 |             inverse(DATASET = args.dataset, network = args.network, NIters = args.iters, imageWidth = imageWidth, inverseClass = c,
231 |             imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer,
232 |             BatchSize = args.batch_size, learningRate = args.learning_rate, NDecreaseLR = args.decrease_LR, eps = args.eps, lambda_TV = args.lambda_TV, lambda_l2 = args.lambda_l2,
233 |             AMSGrad = args.AMSGrad, model_dir = model_dir, model_name = model_name, save_img_dir = save_img_dir, saveIter = args.save_iter,
234 |             gpu = args.gpu, validation=args.validation)
235 | 
236 |     except:
237 |         traceback.print_exc(file=sys.stdout)
238 |         sys.exit(1)
239 | 


--------------------------------------------------------------------------------
/inverse_whitebox_CIFAR_defense.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2020-04-20
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | from net import *
 17 | from utils import *
 18 | 
 19 | from skimage.measure import compare_ssim
 20 | 
 21 | #####################
 22 | # Useful Hyperparameters:
 23 | 
 24 | # CIFAR conv11
 25 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer conv11 --lambda_TV 0.0 --lambda_l2 0.0
 26 | 
 27 | # CIFAR ReLU22
 28 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer ReLU22 --lambda_TV 1e1 --lambda_l2 0.0
 29 | 
 30 | # CIFAR ReLU32
 31 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-2 --layer ReLU32 --lambda_TV 1e3 --lambda_l2 0.0
 32 | 
 33 | # CIFAR fc1
 34 | # python inverse_whitebox_CIFAR.py --iters 5000 --learning_rate 1e-3 --layer fc1 --lambda_TV 1e4 --lambda_l2 0.0
 35 | # python inverse_whitebox_CIFAR.py --iters 1000 --learning_rate 5e-2 --layer fc1 --lambda_TV 1e5 --lambda_l2 0.0
 36 | 
 37 | # CIFAR fc2
 38 | # python inverse_whitebox_CIFAR.py --iters 500 --learning_rate 1e-1 --layer fc2 --lambda_TV 5e2 --lambda_l2 1e2
 39 | 
 40 | # CIFAR label
 41 | # python inverse_whitebox_CIFAR.py --iters 500 --learning_rate 5e-2 --layer label --lambda_TV 5e-1 --lambda_l2 0.0
 42 | 
 43 | # Gaussian and laplace noise
 44 | # python inverse_whitebox_CIFAR_defense.py --noise_type Laplace --layer ReLU22
 45 | # python inverse_whitebox_CIFAR_defense.py --noise_type Laplace --layer ReLU22 --add_noise_to_input
 46 | 
 47 | # Dropout
 48 | # python inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer ReLU22
 49 | # python inverse_whitebox_CIFAR_defense.py --noise_type dropout --layer ReLU22 --add_noise_to_input
 50 | 
 51 | #####################
 52 | 
 53 | def eval_DP_defense(args, noise_type, noise_level, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth"):
 54 |     if args.dataset == 'CIFAR10':
 55 | 
 56 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
 57 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
 58 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 59 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 60 | 
 61 |         tsf = {
 62 |             'train': transforms.Compose(
 63 |             [
 64 |             transforms.ToTensor(),
 65 |             Normalize
 66 |             ]),
 67 |             'test': transforms.Compose(
 68 |             [
 69 |             transforms.ToTensor(),
 70 |             Normalize
 71 |             ])
 72 |         }
 73 | 
 74 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
 75 |                                        download=False, transform = tsf['test'])
 76 | 
 77 |     #print "len(testset) ", len(testset)
 78 |     x_test, y_test = testset.data, testset.targets,
 79 | 
 80 |     #print "x_test.shape ", x_test.shape
 81 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 82 |                                       shuffle = False, num_workers = 1)
 83 |     testIter = iter(testloader)
 84 | 
 85 |     net = torch.load(model_dir + model_name)
 86 |     if not args.gpu:
 87 |         net = net.cpu()
 88 | 
 89 |     net.eval()
 90 |     #print "Validate the model accuracy..."
 91 | 
 92 |     acc = evalTestSplitModel(testloader, net, net, layer=args.layer, gpu = args.gpu,
 93 |             noise_type = noise_type,
 94 |             noise_level = noise_level,
 95 |             args = args
 96 |         )
 97 |     return acc
 98 | 
 99 | def inverse(DATASET = 'CIFAR10', network = 'CIFAR10CNN', NIters = 500, imageWidth = 32, inverseClass = None,
100 |         imageHeight = 32, imageSize = 32*32, NChannels = 3, NClasses = 10, layer = 'conv22',
101 |         BatchSize = 32, learningRate = 1e-3, NDecreaseLR = 20, eps = 1e-3, lambda_TV = 1e3, lambda_l2 = 1.0,
102 |         AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth",
103 |         save_img_dir = "inverted/CIFAR10/MSE_TV/", saveIter = 10, gpu = True, validation=False,
104 |         noise_type = None, noise_level = 0.0, args=None):
105 | 
106 |     assert inverseClass < NClasses
107 | 
108 |     if DATASET == 'CIFAR10':
109 | 
110 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
111 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
112 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
113 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
114 | 
115 |         tsf = {
116 |             'train': transforms.Compose(
117 |             [
118 |             transforms.ToTensor(),
119 |             Normalize
120 |             ]),
121 |             'test': transforms.Compose(
122 |             [
123 |             transforms.ToTensor(),
124 |             Normalize
125 |             ])
126 |         }
127 | 
128 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
129 |                                         download=False, transform = tsf['train'])
130 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
131 |                                        download=False, transform = tsf['test'])
132 | 
133 |     x_train, y_train = trainset.data, trainset.targets,
134 |     x_test, y_test = testset.data, testset.targets,
135 | 
136 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
137 |                                       shuffle = False, num_workers = 1)
138 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
139 |                                       shuffle = False, num_workers = 1)
140 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
141 |                                       shuffle = False, num_workers = 1)
142 |     trainIter = iter(trainloader)
143 |     testIter = iter(testloader)
144 |     inverseIter = iter(inverseloader)
145 | 
146 |     net = torch.load(model_dir + model_name)
147 |     if not gpu:
148 |         net = net.cpu()
149 | 
150 |     net.eval()
151 | 
152 |     targetImg, _ = getImgByClass(inverseIter, C = inverseClass)
153 |     deprocessImg = deprocess(targetImg.clone())
154 | 
155 |     if not os.path.exists(save_img_dir):
156 |         os.makedirs(save_img_dir)
157 |     torchvision.utils.save_image(deprocessImg, save_img_dir + str(inverseClass) + '-ref.png')
158 | 
159 |     if gpu:
160 |         targetImg = targetImg.cuda()
161 |         softmaxLayer = nn.Softmax().cuda()
162 | 
163 |     #if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
164 |     #    targetImg = apply_noise(targetImg, noise_type, noise_level, gpu=args.gpu, args=args)
165 | 
166 |     if layer == 'prob':
167 |         if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
168 |             targetImg_noised = apply_noise(targetImg, noise_type, noise_level, gpu=args.gpu, args=args)
169 |             reflogits = net.forward(targetImg_noised)
170 |         else:
171 |             reflogits = net.forward(targetImg)
172 |         refFeature = softmaxLayer(reflogits)
173 |     elif layer == 'label':
174 |         refFeature = torch.zeros(1,NClasses)
175 |         refFeature[0, inverseClass] = 1
176 |     else:
177 |         targetLayer = net.layerDict[layer]
178 |         if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
179 |             #print "Noise added to input"
180 |             targetImg_noised = apply_noise(targetImg, noise_type, noise_level, gpu=args.gpu, args=args)
181 |             refFeature = net.getLayerOutput(targetImg_noised, targetLayer)
182 |         else:
183 |             refFeature = net.getLayerOutput(targetImg, targetLayer)
184 | 
185 |     # Apply noise
186 |     if noise_type != None and not (hasattr(args, 'add_noise_to_input') and args.add_noise_to_input):
187 |         refFeature = apply_noise(refFeature, noise_type, noise_level, gpu=args.gpu, args=args)
188 | 
189 |     if gpu:
190 |         xGen = torch.zeros(targetImg.size(), requires_grad = True, device="cuda")
191 |     else:
192 |         xGen = torch.zeros(targetImg.size(), requires_grad = True)
193 | 
194 |     optimizer = optim.Adam(params = [xGen], lr = learningRate, eps = eps, amsgrad = AMSGrad)
195 | 
196 |     for i in range(NIters):
197 | 
198 |         optimizer.zero_grad()
199 |         if layer == 'prob':
200 |             xlogits = net.forward(xGen)
201 |             xFeature = softmaxLayer(xlogits)
202 |             featureLoss = ((xFeature - refFeature)**2).mean()
203 |         elif layer == 'label':
204 |             xlogits = net.forward(xGen)
205 |             xFeature = softmaxLayer(xlogits)
206 |             featureLoss = - torch.log(xFeature[0, inverseClass])
207 |         else:
208 |             xFeature = net.getLayerOutput(xGen, targetLayer)
209 |             featureLoss = ((xFeature - refFeature)**2).mean()
210 | 
211 |         TVLoss = TV(xGen)
212 |         normLoss = l2loss(xGen)
213 | 
214 |         totalLoss = featureLoss + lambda_TV * TVLoss + lambda_l2 * normLoss #- 1.0 * conv1Loss
215 | 
216 |         totalLoss.backward(retain_graph=True)
217 |         optimizer.step()
218 | 
219 |         #print "Iter ", i, "Feature loss: ", featureLoss.cpu().detach().numpy(), "TVLoss: ", TVLoss.cpu().detach().numpy(), "l2Loss: ", normLoss.cpu().detach().numpy()
220 | 
221 |     # save the final result
222 |     imgGen = xGen.clone()
223 |     imgGen = deprocess(imgGen)
224 |     torchvision.utils.save_image(imgGen, save_img_dir + str(inverseClass) + '-inv.png')
225 | 
226 |     ref_img = deprocessImg.detach().cpu().numpy().squeeze()
227 |     inv_img = imgGen.detach().cpu().numpy().squeeze()
228 | 
229 |     #print "ref_img.shape", ref_img.shape, "inv_img.shape", inv_img.shape
230 |     #print "ref_img ", ref_img.min(), ref_img.max()
231 |     #print "inv_img ", inv_img.min(), inv_img.max()
232 | 
233 |     psnr = get_PSNR(ref_img, inv_img, peak=1.0)
234 | 
235 |     #print("ref_img.shape", ref_img.shape)
236 |     #print("inv_img.shape", inv_img.shape)
237 |     ssim = compare_ssim(
238 |         X=np.moveaxis(ref_img, 0, -1),
239 |         Y=np.moveaxis(inv_img, 0, -1),
240 |         data_range = inv_img.max() - inv_img.min(),
241 |         multichannel=True)
242 | 
243 |     #print "targetImg l1 Stat:"
244 |     #getL1Stat(net, targetImg)
245 |     #print "xGen l1 Stat:"
246 |     #getL1Stat(net, xGen)
247 |     #print "Done"
248 | 
249 |     return psnr, ssim
250 | 
251 | 
252 | if __name__ == '__main__':
253 |     import argparse
254 |     import sys
255 |     import traceback
256 | 
257 |     try:
258 |         parser = argparse.ArgumentParser()
259 |         parser.add_argument('--dataset', type = str, default = 'CIFAR10')
260 |         parser.add_argument('--network', type = str, default = 'CIFAR10CNN')
261 |         parser.add_argument('--iters', type = int, default = 500)
262 |         parser.add_argument('--eps', type = float, default = 1e-3)
263 |         parser.add_argument('--lambda_TV', type = float, default = 1.0)
264 |         parser.add_argument('--lambda_l2', type = float, default = 1.0)
265 |         parser.add_argument('--AMSGrad', type = bool, default = True)
266 |         parser.add_argument('--batch_size', type = int, default = 32)
267 |         parser.add_argument('--learning_rate', type = float, default = 1e-3)
268 |         parser.add_argument('--decrease_LR', type = int, default = 20)
269 |         parser.add_argument('--layer', type = str, default = 'conv22')
270 |         parser.add_argument('--save_iter', type = int, default = 10)
271 |         parser.add_argument('--inverseClass', type = int, default = None)
272 | 
273 |         parser.add_argument('--noise_type', type = str, default = None)
274 |         parser.add_argument('--noise_level', type = float, default = None)
275 | 
276 |         parser.add_argument('--add_noise_to_input', dest='add_noise_to_input', action='store_true')
277 | 
278 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
279 |         parser.set_defaults(gpu=True)
280 | 
281 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
282 |         parser.set_defaults(validation=True)
283 | 
284 |         parser.add_argument('--noise_iters', type = int, default = 500)
285 |         parser.add_argument('--noise_eps', type = float, default = 1e-3)
286 |         parser.add_argument('--noise_AMSGrad', type = bool, default = True)
287 |         parser.add_argument('--noise_learning_rate', type = float, default = 1e-1)
288 |         parser.add_argument('--noise_lambda_sourcelayer', type = float, default = 1e-1)
289 |         parser.add_argument('--noise_decrease_LR', type = int, default = 20)
290 |         parser.add_argument('--noise_targetLayer', type = str, default = 'fc3')
291 | 
292 |         args = parser.parse_args()
293 | 
294 | 
295 |         args.noise_sourceLayer = args.layer
296 |         args.model_dir = "checkpoints/" + args.dataset + '/'
297 |         args.model_name = "ckpt.pth"
298 | 
299 |         if args.dataset == 'CIFAR10':
300 | 
301 |             imageWidth = 32
302 |             imageHeight = 32
303 |             imageSize = imageWidth * imageHeight
304 |             NChannels = 3
305 |             NClasses = 10
306 | 
307 |         else:
308 |             print "No Dataset Found"
309 |             exit()
310 | 
311 |         noise_type = args.noise_type
312 |         noise_hist = []
313 |         acc_hist = []
314 |         psnr_hist = []
315 |         ssim_hist = []
316 | 
317 |         if 'noise_gen' in args.noise_type:
318 |             default_nl = np.concatenate((np.arange(0, 110, 10), np.arange(100, 1100, 100)), axis=0)
319 |         elif 'dropout' in args.noise_type:
320 |             default_nl = np.arange(0, 1, 0.1)
321 |         else:
322 |             default_nl = np.concatenate((np.arange(0, 1, 0.1), np.arange(1.0, 5.5, 0.5)), axis=0)
323 | 
324 |         noise_range = [args.noise_level] if args.noise_level != None else default_nl
325 | 
326 |         for noise_level in noise_range:
327 |             noise_hist.append(noise_level)
328 | 
329 |             if args.add_noise_to_input:
330 |                 save_img_dir = "inverted_whitebox/" + args.dataset + '/' + args.layer + '/' + 'noised_add_to_input/' + noise_type + '/' + str(round(noise_level,1)) + '/'
331 |             else:
332 |                 save_img_dir = "inverted_whitebox/" + args.dataset + '/' + args.layer + '/' + 'noised/' + noise_type + '/' + str(round(noise_level,1)) + '/'
333 | 
334 |             if not os.path.exists(save_img_dir):
335 |                 os.makedirs(save_img_dir)
336 | 
337 |             acc = eval_DP_defense(args, noise_type, noise_level)
338 |             acc_hist.append(acc)
339 | 
340 |             psnr_sum = 0.0
341 |             ssim_sum = 0.0
342 |             for c in range(NClasses):
343 |                 psnr, ssim = inverse(DATASET = args.dataset, network = args.network, NIters = args.iters, imageWidth = imageWidth, inverseClass = c,
344 |                 imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer,
345 |                 BatchSize = args.batch_size, learningRate = args.learning_rate, NDecreaseLR = args.decrease_LR, eps = args.eps, lambda_TV = args.lambda_TV, lambda_l2 = args.lambda_l2,
346 |                 AMSGrad = args.AMSGrad, model_dir = args.model_dir, model_name = args.model_name, save_img_dir = save_img_dir, saveIter = args.save_iter,
347 |                 gpu = args.gpu, validation=args.validation, noise_type = noise_type, noise_level = noise_level, args = args)
348 | 
349 |                 psnr_sum += psnr / NClasses
350 |                 ssim_sum += ssim / NClasses
351 | 
352 |             psnr_hist.append(psnr_sum)
353 |             ssim_hist.append(ssim_sum)
354 |             print "Noise_type:", noise_type, " Add to input:", args.add_noise_to_input, " Noise_level:", round(noise_level,2), " Acc:", round(acc,4), " PSNR:", round(psnr_sum,4), " SSIM:", round(ssim_sum,4)
355 | 
356 |     except:
357 |         traceback.print_exc(file=sys.stdout)
358 |         sys.exit(1)
359 | 


--------------------------------------------------------------------------------
/inverse_whitebox_MNIST.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2020-04-20
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | from net import *
 17 | from utils import *
 18 | 
 19 | #####################
 20 | # Useful Hyperparameters:
 21 | 
 22 | # A good parameter set to invert MNIST ReLU2 layer
 23 | # python inverse.py --dataset MNIST --iters 5000 --learning_rate 1e-2 --layer ReLU2 --lambda_TV 1e0 --lambda_l2 0.0
 24 | 
 25 | # A good parameter set to invert MNIST fc3 layer
 26 | # python inverse.py --dataset MNIST --iters 5000 --learning_rate 1e-2 --layer fc3 --lambda_TV 1e3 --lambda_l2 0.0
 27 | #
 28 | # A good parameter set to invert MNIST label only
 29 | # python inverse.py --dataset MNIST --iters 5000 --learning_rate 1e-2 --layer prob --lambda_TV 1e-1 --lambda_l2 0.0
 30 | 
 31 | #####################
 32 | 
 33 | def inverse(DATASET = 'MNIST', network = 'LeNet', NIters = 500, imageWidth = 28, inverseClass = None,
 34 |         imageHeight = 28, imageSize = 28*28, NChannels = 1, NClasses = 10, layer = 'conv2',
 35 |         BatchSize = 32, learningRate = 1e-3, NDecreaseLR = 20, eps = 1e-3, lambda_TV = 1e3, lambda_l2 = 1.0,
 36 |         AMSGrad = True, model_dir = "checkpoints/MNIST/", model_name = "ckpt.pth",
 37 |         save_img_dir = "inverted/MNIST/MSE_TV/", saveIter = 10, gpu = True, validation=False):
 38 | 
 39 |     print "DATASET: ", DATASET
 40 |     print "inverseClass: ", inverseClass
 41 | 
 42 |     assert inverseClass < NClasses
 43 | 
 44 |     if DATASET == 'MNIST':
 45 | 
 46 |         mu = torch.tensor([0.5], dtype=torch.float32)
 47 |         sigma = torch.tensor([0.5], dtype=torch.float32)
 48 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 49 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 50 | 
 51 |         tsf = {
 52 |             'train': transforms.Compose(
 53 |             [
 54 |                 transforms.ToTensor(),
 55 |                 Normalize
 56 |             ]),
 57 | 
 58 |             'test': transforms.Compose(
 59 |             [
 60 |                 transforms.ToTensor(),
 61 |                 Normalize
 62 |             ])
 63 |         }
 64 | 
 65 |         trainset = torchvision.datasets.MNIST(root='./data/MNIST', train=True,
 66 |                                         download=True, transform = tsf['train'])
 67 | 
 68 |         testset = torchvision.datasets.MNIST(root='./data/MNIST', train=False,
 69 |                                        download=True, transform = tsf['test'])
 70 | 
 71 |     print "len(trainset) ", len(trainset)
 72 |     print "len(testset) ", len(testset)
 73 |     x_train, y_train = trainset.data, trainset.targets,
 74 |     x_test, y_test = testset.data, testset.targets,
 75 | 
 76 |     print "x_train.shape ", x_train.shape
 77 |     print "x_test.shape ", x_test.shape
 78 | 
 79 | 
 80 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
 81 |                                       shuffle = False, num_workers = 1)
 82 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 83 |                                       shuffle = False, num_workers = 1)
 84 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
 85 |                                       shuffle = False, num_workers = 1)
 86 |     trainIter = iter(trainloader)
 87 |     testIter = iter(testloader)
 88 |     inverseIter = iter(inverseloader)
 89 | 
 90 |     net = torch.load(model_dir + model_name)
 91 |     if not gpu:
 92 |         net = net.cpu()
 93 | 
 94 |     net.eval()
 95 |     print "Validate the model accuracy..."
 96 |     if validation:
 97 |         accTest = evalTest(testloader, net, gpu = gpu)
 98 | 
 99 |     targetImg, _ = getImgByClass(inverseIter, C = inverseClass)
100 |     print "targetImg.size()", targetImg.size()
101 | 
102 |     deprocessImg = deprocess(targetImg.clone())
103 | 
104 |     if not os.path.exists(save_img_dir):
105 |         os.makedirs(save_img_dir)
106 |     torchvision.utils.save_image(deprocessImg, save_img_dir + str(inverseClass) + '-ref.png')
107 | 
108 |     if gpu:
109 |         targetImg = targetImg.cuda()
110 |         softmaxLayer = nn.Softmax().cuda()
111 | 
112 |     if layer == 'prob':
113 |         reflogits = net.forward(targetImg)
114 |         refFeature = softmaxLayer(reflogits)
115 |     elif layer == 'label':
116 |         refFeature = torch.zeros(1,NClasses)
117 |         refFeature[0, inverseClass] = 1
118 |     else:
119 |         targetLayer = net.layerDict[layer]
120 |         refFeature = net.getLayerOutput(targetImg, targetLayer)
121 | 
122 |     print "refFeature.size()", refFeature.size()
123 | 
124 |     if gpu:
125 |         xGen = torch.zeros(targetImg.size(), requires_grad = True, device="cuda")
126 |     else:
127 |         xGen = torch.zeros(targetImg.size(), requires_grad = True)
128 | 
129 |     optimizer = optim.Adam(params = [xGen], lr = learningRate, eps = eps, amsgrad = AMSGrad)
130 | 
131 |     for i in range(NIters):
132 | 
133 |         optimizer.zero_grad()
134 |         if layer == 'prob':
135 |             xlogits = net.forward(xGen)
136 |             xFeature = softmaxLayer(xlogits)
137 |             featureLoss = ((xFeature - refFeature)**2).mean()
138 |         elif layer == 'label':
139 |             xlogits = net.forward(xGen)
140 |             xFeature = softmaxLayer(xlogits)
141 |             featureLoss = - torch.log(xFeature[0, inverseClass])
142 |         else:
143 |             xFeature = net.getLayerOutput(xGen, targetLayer)
144 |             featureLoss = ((xFeature - refFeature)**2).mean()
145 | 
146 |         TVLoss = TV(xGen)
147 |         normLoss = l2loss(xGen)
148 | 
149 |         totalLoss = featureLoss + lambda_TV * TVLoss + lambda_l2 * normLoss #- 1.0 * conv1Loss
150 | 
151 |         totalLoss.backward(retain_graph=True)
152 |         optimizer.step()
153 | 
154 |         print "Iter ", i, "Feature loss: ", featureLoss.cpu().detach().numpy(), "TVLoss: ", TVLoss.cpu().detach().numpy(), "l2Loss: ", normLoss.cpu().detach().numpy()
155 | 
156 |     # save the final result
157 |     imgGen = xGen.clone()
158 |     imgGen = deprocess(imgGen)
159 |     torchvision.utils.save_image(imgGen, save_img_dir + str(inverseClass) + '-inv.png')
160 | 
161 |     print "targetImg l1 Stat:"
162 |     getL1Stat(net, targetImg)
163 |     print "xGen l1 Stat:"
164 |     getL1Stat(net, xGen)
165 |     print "Done"
166 | 
167 | 
168 | if __name__ == '__main__':
169 |     import argparse
170 |     import sys
171 |     import traceback
172 | 
173 |     try:
174 |         parser = argparse.ArgumentParser()
175 |         parser.add_argument('--dataset', type = str, default = 'MNIST')
176 |         parser.add_argument('--network', type = str, default = 'LeNet')
177 |         parser.add_argument('--iters', type = int, default = 500)
178 |         parser.add_argument('--eps', type = float, default = 1e-3)
179 |         parser.add_argument('--lambda_TV', type = float, default = 1.0)
180 |         parser.add_argument('--lambda_l2', type = float, default = 0.0)
181 |         parser.add_argument('--AMSGrad', type = bool, default = True)
182 |         parser.add_argument('--batch_size', type = int, default = 32)
183 |         parser.add_argument('--learning_rate', type = float, default = 1e-2)
184 |         parser.add_argument('--decrease_LR', type = int, default = 20)
185 |         parser.add_argument('--layer', type = str, default = 'ReLU2')
186 |         parser.add_argument('--save_iter', type = int, default = 10)
187 |         parser.add_argument('--inverseClass', type = int, default = None)
188 | 
189 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
190 |         parser.set_defaults(gpu=True)
191 | 
192 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
193 |         parser.set_defaults(validation=True)
194 |         args = parser.parse_args()
195 | 
196 |         model_dir = "checkpoints/" + args.dataset + '/'
197 |         model_name = "ckpt.pth"
198 | 
199 |         save_img_dir = "inverted_whitebox/" + args.dataset + '/' + args.layer + '/'
200 | 
201 |         if args.dataset == 'MNIST':
202 | 
203 |             imageWidth = 28
204 |             imageHeight = 28
205 |             imageSize = imageWidth * imageHeight
206 |             NChannels = 1
207 |             NClasses = 10
208 | 
209 |         else:
210 |             print "No Dataset Found"
211 |             exit()
212 | 
213 |         for c in range(NClasses):
214 |             inverse(DATASET = args.dataset, network = args.network, NIters = args.iters, imageWidth = imageWidth, inverseClass = c,
215 |             imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer,
216 |             BatchSize = args.batch_size, learningRate = args.learning_rate, NDecreaseLR = args.decrease_LR, eps = args.eps, lambda_TV = args.lambda_TV, lambda_l2 = args.lambda_l2,
217 |             AMSGrad = args.AMSGrad, model_dir = model_dir, model_name = model_name, save_img_dir = save_img_dir, saveIter = args.save_iter,
218 |             gpu = args.gpu, validation=args.validation)
219 | 
220 |     except:
221 |         traceback.print_exc(file=sys.stdout)
222 |         sys.exit(1)
223 | 


--------------------------------------------------------------------------------
/inverse_whitebox_MNIST_defense.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2020-04-20
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | from net import *
 17 | from utils import *
 18 | 
 19 | from skimage.measure import compare_ssim
 20 | 
 21 | #####################
 22 | # Useful Hyperparameters:
 23 | 
 24 | # A good parameter set to invert MNIST ReLU2 layer
 25 | # python inverse.py --dataset MNIST --iters 5000 --learning_rate 1e-2 --layer ReLU2 --lambda_TV 1e0 --lambda_l2 0.0
 26 | 
 27 | # A good parameter set to invert MNIST fc3 layer
 28 | # python inverse.py --dataset MNIST --iters 5000 --learning_rate 1e-2 --layer fc3 --lambda_TV 1e3 --lambda_l2 0.0
 29 | #
 30 | # A good parameter set to invert MNIST label only
 31 | # python inverse.py --dataset MNIST --iters 5000 --learning_rate 1e-2 --layer prob --lambda_TV 1e-1 --lambda_l2 0.0
 32 | 
 33 | # Gaussian and laplace noise
 34 | # python inverse_whitebox_MNIST_defense.py --noise_type Laplace --layer ReLU2
 35 | # python inverse_whitebox_MNIST_defense.py --noise_type Laplace --layer ReLU2 --add_noise_to_input
 36 | 
 37 | # Dropout
 38 | # python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer ReLU2
 39 | # python inverse_whitebox_MNIST_defense.py --noise_type dropout --layer ReLU2 --add_noise_to_input
 40 | 
 41 | # Noise opt
 42 | # Generate a noise applied to noise_targetLayer that minimize the difference in noise_targetLayer
 43 | # python noise_generation_opt.py --noise_sourceLayer ReLU2 --noise_targetLayer fc1
 44 | # python inverse_whitebox_MNIST_defense.py --noise_type noise_gen --noise_targetLayer fc1
 45 | # python noise_generation_opt.py --noise_sourceLayer conv2 --noise_targetLayer ReLU2 --noise_lambda_sourcelayer 1e-3 --noise_level 200
 46 | 
 47 | #####################
 48 | 
 49 | def eval_DP_defense(args, noise_type, noise_level, model_dir = "checkpoints/MNIST/", model_name = "ckpt.pth"):
 50 |     if args.dataset == 'MNIST':
 51 | 
 52 |         mu = torch.tensor([0.5], dtype=torch.float32)
 53 |         sigma = torch.tensor([0.5], dtype=torch.float32)
 54 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 55 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 56 | 
 57 |         tsf = {
 58 |             'test': transforms.Compose(
 59 |             [
 60 |                 transforms.ToTensor(),
 61 |                 Normalize
 62 |             ])
 63 |         }
 64 | 
 65 |         testset = torchvision.datasets.MNIST(root='./data/MNIST', train=False,
 66 |                                        download=True, transform = tsf['test'])
 67 | 
 68 |     #print "len(testset) ", len(testset)
 69 |     x_test, y_test = testset.data, testset.targets,
 70 | 
 71 |     #print "x_test.shape ", x_test.shape
 72 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 73 |                                       shuffle = False, num_workers = 1)
 74 |     testIter = iter(testloader)
 75 | 
 76 |     net = torch.load(model_dir + model_name)
 77 |     if not args.gpu:
 78 |         net = net.cpu()
 79 | 
 80 |     net.eval()
 81 |     #print "Validate the model accuracy..."
 82 | 
 83 |     acc = evalTestSplitModel(testloader, net, net, layer=args.layer, gpu = args.gpu,
 84 |             noise_type = noise_type,
 85 |             noise_level = noise_level,
 86 |             args = args
 87 |         )
 88 |     return acc
 89 | 
 90 | def inverse(DATASET = 'MNIST', network = 'LeNet', NIters = 500, imageWidth = 28, inverseClass = None,
 91 |         imageHeight = 28, imageSize = 28*28, NChannels = 1, NClasses = 10, layer = 'conv2',
 92 |         BatchSize = 32, learningRate = 1e-3, NDecreaseLR = 20, eps = 1e-3, lambda_TV = 1e3, lambda_l2 = 1.0,
 93 |         AMSGrad = True, model_dir = "checkpoints/MNIST/", model_name = "ckpt.pth",
 94 |         save_img_dir = "inverted/MNIST/MSE_TV/", saveIter = 10, gpu = True, validation=False,
 95 |         noise_type = None, noise_level = 0.0, args=None):
 96 | 
 97 |     assert inverseClass < NClasses
 98 | 
 99 |     if DATASET == 'MNIST':
100 | 
101 |         mu = torch.tensor([0.5], dtype=torch.float32)
102 |         sigma = torch.tensor([0.5], dtype=torch.float32)
103 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
104 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
105 | 
106 |         tsf = {
107 |             'train': transforms.Compose(
108 |             [
109 |                 transforms.ToTensor(),
110 |                 Normalize
111 |             ]),
112 | 
113 |             'test': transforms.Compose(
114 |             [
115 |                 transforms.ToTensor(),
116 |                 Normalize
117 |             ])
118 |         }
119 | 
120 |         trainset = torchvision.datasets.MNIST(root='./data/MNIST', train=True,
121 |                                         download=True, transform = tsf['train'])
122 | 
123 |         testset = torchvision.datasets.MNIST(root='./data/MNIST', train=False,
124 |                                        download=True, transform = tsf['test'])
125 | 
126 |     x_train, y_train = trainset.data, trainset.targets,
127 |     x_test, y_test = testset.data, testset.targets,
128 | 
129 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
130 |                                       shuffle = False, num_workers = 1)
131 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
132 |                                       shuffle = False, num_workers = 1)
133 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
134 |                                       shuffle = False, num_workers = 1)
135 |     trainIter = iter(trainloader)
136 |     testIter = iter(testloader)
137 |     inverseIter = iter(inverseloader)
138 | 
139 |     net = torch.load(model_dir + model_name)
140 |     if not gpu:
141 |         net = net.cpu()
142 | 
143 |     net.eval()
144 | 
145 |     targetImg, _ = getImgByClass(inverseIter, C = inverseClass)
146 |     deprocessImg = deprocess(targetImg.clone())
147 | 
148 |     if not os.path.exists(save_img_dir):
149 |         os.makedirs(save_img_dir)
150 |     torchvision.utils.save_image(deprocessImg, save_img_dir + str(inverseClass) + '-ref.png')
151 | 
152 |     if gpu:
153 |         targetImg = targetImg.cuda()
154 |         softmaxLayer = nn.Softmax().cuda()
155 | 
156 |     #if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
157 |     #    targetImg = apply_noise(targetImg, noise_type, noise_level, gpu=args.gpu, args=args)
158 | 
159 |     if layer == 'prob':
160 |         if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
161 |             targetImg_noised = apply_noise(targetImg, noise_type, noise_level, gpu=args.gpu, args=args)
162 |             reflogits = net.forward(targetImg_noised)
163 |         else:
164 |             reflogits = net.forward(targetImg)
165 |         refFeature = softmaxLayer(reflogits)
166 |     elif layer == 'label':
167 |         refFeature = torch.zeros(1,NClasses)
168 |         refFeature[0, inverseClass] = 1
169 |     else:
170 |         targetLayer = net.layerDict[layer]
171 |         if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
172 |             #print "Noise added to input"
173 |             targetImg_noised = apply_noise(targetImg, noise_type, noise_level, gpu=args.gpu, args=args)
174 |             refFeature = net.getLayerOutput(targetImg_noised, targetLayer)
175 |         else:
176 |             refFeature = net.getLayerOutput(targetImg, targetLayer)
177 | 
178 |     # Apply noise
179 |     if noise_type != None and not (hasattr(args, 'add_noise_to_input') and args.add_noise_to_input):
180 |         refFeature = apply_noise(refFeature, noise_type, noise_level, gpu=args.gpu, args=args)
181 | 
182 |     if gpu:
183 |         xGen = torch.zeros(targetImg.size(), requires_grad = True, device="cuda")
184 |     else:
185 |         xGen = torch.zeros(targetImg.size(), requires_grad = True)
186 | 
187 |     optimizer = optim.Adam(params = [xGen], lr = learningRate, eps = eps, amsgrad = AMSGrad)
188 | 
189 |     for i in range(NIters):
190 | 
191 |         optimizer.zero_grad()
192 |         if layer == 'prob':
193 |             xlogits = net.forward(xGen)
194 |             xFeature = softmaxLayer(xlogits)
195 |             featureLoss = ((xFeature - refFeature)**2).mean()
196 |         elif layer == 'label':
197 |             xlogits = net.forward(xGen)
198 |             xFeature = softmaxLayer(xlogits)
199 |             featureLoss = - torch.log(xFeature[0, inverseClass])
200 |         else:
201 |             xFeature = net.getLayerOutput(xGen, targetLayer)
202 |             featureLoss = ((xFeature - refFeature)**2).mean()
203 | 
204 |         TVLoss = TV(xGen)
205 |         normLoss = l2loss(xGen)
206 | 
207 |         totalLoss = featureLoss + lambda_TV * TVLoss + lambda_l2 * normLoss #- 1.0 * conv1Loss
208 | 
209 |         totalLoss.backward(retain_graph=True)
210 |         optimizer.step()
211 | 
212 |         #print "Iter ", i, "Feature loss: ", featureLoss.cpu().detach().numpy(), "TVLoss: ", TVLoss.cpu().detach().numpy(), "l2Loss: ", normLoss.cpu().detach().numpy()
213 | 
214 |     # save the final result
215 |     imgGen = xGen.clone()
216 |     imgGen = deprocess(imgGen)
217 |     torchvision.utils.save_image(imgGen, save_img_dir + str(inverseClass) + '-inv.png')
218 | 
219 |     ref_img = deprocessImg.detach().cpu().numpy().squeeze()
220 |     inv_img = imgGen.detach().cpu().numpy().squeeze()
221 | 
222 |     #print "ref_img.shape", ref_img.shape, "inv_img.shape", inv_img.shape
223 |     #print "ref_img ", ref_img.min(), ref_img.max()
224 |     #print "inv_img ", inv_img.min(), inv_img.max()
225 | 
226 |     psnr = get_PSNR(ref_img, inv_img, peak=1.0)
227 |     ssim = compare_ssim(ref_img, inv_img, data_range = inv_img.max() - inv_img.min(), multichannel=False)
228 | 
229 |     #print "targetImg l1 Stat:"
230 |     #getL1Stat(net, targetImg)
231 |     #print "xGen l1 Stat:"
232 |     #getL1Stat(net, xGen)
233 |     #print "Done"
234 | 
235 |     return psnr, ssim
236 | 
237 | 
238 | if __name__ == '__main__':
239 |     import argparse
240 |     import sys
241 |     import traceback
242 | 
243 |     try:
244 |         parser = argparse.ArgumentParser()
245 |         parser.add_argument('--dataset', type = str, default = 'MNIST')
246 |         parser.add_argument('--network', type = str, default = 'LeNet')
247 |         parser.add_argument('--iters', type = int, default = 500)
248 |         parser.add_argument('--eps', type = float, default = 1e-3)
249 |         parser.add_argument('--lambda_TV', type = float, default = 1.0)
250 |         parser.add_argument('--lambda_l2', type = float, default = 0.0)
251 |         parser.add_argument('--AMSGrad', type = bool, default = True)
252 |         parser.add_argument('--batch_size', type = int, default = 32)
253 |         parser.add_argument('--learning_rate', type = float, default = 1e-2)
254 |         parser.add_argument('--decrease_LR', type = int, default = 20)
255 |         parser.add_argument('--layer', type = str, default = 'ReLU2')
256 |         parser.add_argument('--save_iter', type = int, default = 10)
257 |         parser.add_argument('--inverseClass', type = int, default = None)
258 |         parser.add_argument('--noise_type', type = str, default = None)
259 |         parser.add_argument('--noise_level', type = float, default = None)
260 | 
261 |         parser.add_argument('--add_noise_to_input', dest='add_noise_to_input', action='store_true')
262 | 
263 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
264 |         parser.set_defaults(gpu=True)
265 | 
266 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
267 |         parser.set_defaults(validation=True)
268 | 
269 |         parser.add_argument('--noise_iters', type = int, default = 500)
270 |         parser.add_argument('--noise_eps', type = float, default = 1e-3)
271 |         parser.add_argument('--noise_AMSGrad', type = bool, default = True)
272 |         parser.add_argument('--noise_learning_rate', type = float, default = 1e-1)
273 |         parser.add_argument('--noise_lambda_sourcelayer', type = float, default = 1e-1)
274 |         parser.add_argument('--noise_decrease_LR', type = int, default = 20)
275 |         parser.add_argument('--noise_targetLayer', type = str, default = 'fc3')
276 | 
277 |         args = parser.parse_args()
278 | 
279 | 
280 |         args.noise_sourceLayer = args.layer
281 |         args.model_dir = "checkpoints/" + args.dataset + '/'
282 |         args.model_name = "ckpt.pth"
283 | 
284 |         if args.dataset == 'MNIST':
285 | 
286 |             imageWidth = 28
287 |             imageHeight = 28
288 |             imageSize = imageWidth * imageHeight
289 |             NChannels = 1
290 |             NClasses = 10
291 | 
292 |         else:
293 |             print "No Dataset Found"
294 |             exit()
295 | 
296 |         noise_type = args.noise_type
297 |         noise_hist = []
298 |         acc_hist = []
299 |         psnr_hist = []
300 |         ssim_hist = []
301 | 
302 |         if 'noise_gen' in args.noise_type:
303 |             default_nl = np.concatenate((np.arange(0, 110, 10), np.arange(100, 1100, 100)), axis=0)
304 |         elif 'dropout' in args.noise_type:
305 |             default_nl = np.arange(0, 1, 0.1)
306 |         else:
307 |             default_nl = np.concatenate((np.arange(0, 1, 0.1), np.arange(1.0, 5.5, 0.5)), axis=0)
308 | 
309 |         noise_range = [args.noise_level] if args.noise_level != None else default_nl
310 | 
311 |         for noise_level in noise_range:
312 |             noise_hist.append(noise_level)
313 | 
314 |             if args.add_noise_to_input:
315 |                 save_img_dir = "inverted_whitebox/" + args.dataset + '/' + args.layer + '/' + 'noised_add_to_input/' + noise_type + '/' + str(round(noise_level,1)) + '/'
316 |             else:
317 |                 save_img_dir = "inverted_whitebox/" + args.dataset + '/' + args.layer + '/' + 'noised/' + noise_type + '/' + str(round(noise_level,1)) + '/'
318 | 
319 |             if not os.path.exists(save_img_dir):
320 |                 os.makedirs(save_img_dir)
321 | 
322 |             acc = eval_DP_defense(args, noise_type, noise_level)
323 |             acc_hist.append(acc)
324 | 
325 |             psnr_sum = 0.0
326 |             ssim_sum = 0.0
327 |             for c in range(NClasses):
328 |                 psnr, ssim = inverse(DATASET = args.dataset, network = args.network, NIters = args.iters, imageWidth = imageWidth, inverseClass = c,
329 |                 imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses, layer = args.layer,
330 |                 BatchSize = args.batch_size, learningRate = args.learning_rate, NDecreaseLR = args.decrease_LR, eps = args.eps, lambda_TV = args.lambda_TV, lambda_l2 = args.lambda_l2,
331 |                 AMSGrad = args.AMSGrad, model_dir = args.model_dir, model_name = args.model_name, save_img_dir = save_img_dir, saveIter = args.save_iter,
332 |                 gpu = args.gpu, validation=args.validation, noise_type = noise_type, noise_level = noise_level, args = args)
333 | 
334 |                 psnr_sum += psnr / NClasses
335 |                 ssim_sum += ssim / NClasses
336 | 
337 |             psnr_hist.append(psnr_sum)
338 |             ssim_hist.append(ssim_sum)
339 |             print "Noise_type:", noise_type, " Add to input:", args.add_noise_to_input, " Noise_level:", round(noise_level,2), " Acc:", round(acc,4), " PSNR:", round(psnr_sum,4), " SSIM:", round(ssim_sum,4)
340 | 
341 |     except:
342 |         traceback.print_exc(file=sys.stdout)
343 |         sys.exit(1)
344 | 


--------------------------------------------------------------------------------
/net.py:
--------------------------------------------------------------------------------
   1 | # @Author: zechenghe
   2 | # @Date:   2019-01-21T12:17:31-05:00
   3 | # @Last modified by:   zechenghe
   4 | # @Last modified time: 2019-02-01T14:01:15-05:00
   5 | 
   6 | import time
   7 | import math
   8 | import os
   9 | import sys
  10 | import numpy as np
  11 | 
  12 | import torch
  13 | import torchvision
  14 | import torch.nn as nn
  15 | import torch.nn.functional as F
  16 | 
  17 | import collections
  18 | 
  19 | class LeNet(nn.Module):
  20 |     def __init__(self, NChannels):
  21 |         super(LeNet, self).__init__()
  22 |         self.features = []
  23 |         self.classifier = []
  24 |         self.layerDict = collections.OrderedDict()
  25 | 
  26 |         self.conv1 = nn.Conv2d(
  27 |             in_channels = NChannels,
  28 |             out_channels = 8,
  29 |             kernel_size = 5
  30 |         )
  31 |         self.features.append(self.conv1)
  32 |         self.layerDict['conv1'] = self.conv1
  33 | 
  34 |         self.ReLU1 = nn.ReLU(True)
  35 |         self.features.append(self.ReLU1)
  36 |         self.layerDict['ReLU1'] = self.ReLU1
  37 | 
  38 |         self.pool1 = nn.MaxPool2d(2,2)
  39 |         self.features.append(self.pool1)
  40 |         self.layerDict['pool1'] = self.pool1
  41 | 
  42 |         self.conv2 = nn.Conv2d(
  43 |             in_channels = 8,
  44 |             out_channels = 16,
  45 |             kernel_size = 5
  46 |         )
  47 |         self.features.append(self.conv2)
  48 |         self.layerDict['conv2'] = self.conv2
  49 | 
  50 |         self.ReLU2 = nn.ReLU(True)
  51 |         self.features.append(self.ReLU2)
  52 |         self.layerDict['ReLU2'] = self.ReLU2
  53 | 
  54 |         self.pool2 = nn.MaxPool2d(2,2)
  55 |         self.features.append(self.pool2)
  56 |         self.layerDict['pool2'] = self.pool2
  57 | 
  58 |         self.feature_dims = 16 * 4 * 4
  59 |         self.fc1 = nn.Linear(self.feature_dims, 120)
  60 |         self.classifier.append(self.fc1)
  61 |         self.layerDict['fc1'] = self.fc1
  62 | 
  63 |         self.fc1act = nn.ReLU(True)
  64 |         self.classifier.append(self.fc1act)
  65 |         self.layerDict['fc1act'] = self.fc1act
  66 | 
  67 | 
  68 |         self.fc2 = nn.Linear(120, 84)
  69 |         self.classifier.append(self.fc2)
  70 |         self.layerDict['fc2'] = self.fc2
  71 | 
  72 |         self.fc2act = nn.ReLU(True)
  73 |         self.classifier.append(self.fc2act)
  74 |         self.layerDict['fc2act'] = self.fc2act
  75 | 
  76 |         self.fc3 = nn.Linear(84, 10)
  77 |         self.classifier.append(self.fc3)
  78 |         self.layerDict['fc3'] = self.fc3
  79 | 
  80 | 
  81 |     def forward(self, x):
  82 |         for layer in self.features:
  83 |             x = layer(x)
  84 | 
  85 |         #print 'x.size', x.size()
  86 |         x = x.view(-1, self.feature_dims)
  87 | 
  88 |         for layer in self.classifier:
  89 |             x = layer(x)
  90 |         return x
  91 | 
  92 |     def forward_from(self, x, layer):
  93 | 
  94 |         if layer == 'input':
  95 |             return x
  96 | 
  97 |         if layer in self.layerDict:
  98 |             targetLayer = self.layerDict[layer]
  99 | 
 100 |             if targetLayer in self.features:
 101 |                 layeridx = self.features.index(targetLayer)
 102 |                 for func in self.features[layeridx+1:]:
 103 |                     x = func(x)
 104 | #                    print "x.size() ", x.size()
 105 | 
 106 | #                print "Pass Features "
 107 |                 x = x.view(-1, self.feature_dims)
 108 |                 for func in self.classifier:
 109 |                     x = func(x)
 110 |                 return x
 111 | 
 112 |             else:
 113 |                 layeridx = self.classifier.index(targetLayer)
 114 |                 for func in self.classifier[layeridx+1:]:
 115 |                     x = func(x)
 116 |                 return x
 117 |         else:
 118 |             print "layer not exists"
 119 |             exit(1)
 120 | 
 121 | 
 122 |     def getLayerOutput(self, x, targetLayer):
 123 |         if targetLayer == 'input':
 124 |             return x
 125 | 
 126 |         for layer in self.features:
 127 |             x = layer(x)
 128 |             if layer == targetLayer:
 129 |                 return x
 130 | 
 131 |         x = x.view(-1, self.feature_dims)
 132 |         for layer in self.classifier:
 133 |             x = layer(x)
 134 |             if layer == targetLayer:
 135 |                 return x
 136 | 
 137 |         # Should not go here
 138 |         raise Exception("Target layer not found")
 139 | 
 140 |     def getLayerOutputFrom(self, x, sourceLayer, targetLayer):
 141 | 
 142 |         if targetLayer == 'input':
 143 |             return x
 144 | 
 145 |         if sourceLayer == 'input':
 146 |             return getLayerOutput(self, x, self.layerDict[targetLayer])
 147 | 
 148 |         if sourceLayer in self.layerDict and targetLayer in self.layerDict:
 149 |             sourceLayer = self.layerDict[sourceLayer]
 150 |             targetLayer = self.layerDict[targetLayer]
 151 | 
 152 |             if sourceLayer in self.features and targetLayer in self.features:
 153 |                 sourceLayeridx = self.features.index(sourceLayer)
 154 |                 targetLayeridx = self.features.index(targetLayer)
 155 | 
 156 |                 for func in self.features[sourceLayeridx+1:targetLayeridx+1]:
 157 |                     x = func(x)
 158 |                 return x
 159 | 
 160 |             elif sourceLayer in self.classifier and targetLayer in self.classifier:
 161 |                 sourceLayeridx = self.classifier.index(sourceLayer)
 162 |                 targetLayeridx = self.classifier.index(targetLayer)
 163 | 
 164 |                 for func in self.classifier[sourceLayeridx+1:targetLayeridx+1]:
 165 |                     x = func(x)
 166 |                 return x
 167 | 
 168 |             elif sourceLayer in self.features and targetLayer in self.classifier:
 169 |                 sourceLayeridx = self.features.index(sourceLayer)
 170 |                 for func in self.features[sourceLayeridx+1:]:
 171 |                     x = func(x)
 172 | 
 173 |                 x = x.view(-1, self.feature_dims)
 174 |                 targetLayeridx = self.classifier.index(targetLayer)
 175 |                 for func in self.classifier[:targetLayeridx+1]:
 176 |                     x = func(x)
 177 |                 return x
 178 | 
 179 |             else:
 180 |                 print "Target layer cannot before source layer"
 181 |                 exit(1)
 182 |         else:
 183 |             print "layer not exists"
 184 |             exit(1)
 185 | 
 186 | 
 187 | class CIFAR10CNN(nn.Module):
 188 |     def __init__(self, NChannels):
 189 |         super(CIFAR10CNN, self).__init__()
 190 |         self.features = []
 191 |         self.layerDict = collections.OrderedDict()
 192 | 
 193 |         self.conv11 = nn.Conv2d(
 194 |             in_channels = NChannels,
 195 |             out_channels = 64,
 196 |             kernel_size = 3,
 197 |             padding = 1
 198 |         )
 199 |         self.features.append(self.conv11)
 200 |         self.layerDict['conv11'] = self.conv11
 201 | 
 202 |         self.ReLU11 = nn.ReLU()
 203 |         self.features.append(self.ReLU11)
 204 |         self.layerDict['ReLU11'] = self.ReLU11
 205 | 
 206 |         self.conv12 = nn.Conv2d(
 207 |             in_channels = 64,
 208 |             out_channels = 64,
 209 |             kernel_size = 3,
 210 |             padding = 1
 211 |         )
 212 |         self.features.append(self.conv12)
 213 |         self.layerDict['conv12'] = self.conv12
 214 | 
 215 |         self.ReLU12 = nn.ReLU()
 216 |         self.features.append(self.ReLU12)
 217 |         self.layerDict['ReLU12'] = self.ReLU12
 218 | 
 219 |         self.pool1 = nn.MaxPool2d(2,2)
 220 |         self.features.append(self.pool1)
 221 |         self.layerDict['pool1'] = self.pool1
 222 | 
 223 | 
 224 |         self.conv21 = nn.Conv2d(
 225 |             in_channels = 64,
 226 |             out_channels = 128,
 227 |             kernel_size = 3,
 228 |             padding = 1
 229 |         )
 230 |         self.features.append(self.conv21)
 231 |         self.layerDict['conv21'] = self.conv21
 232 | 
 233 |         self.ReLU21 = nn.ReLU()
 234 |         self.features.append(self.ReLU21)
 235 |         self.layerDict['ReLU21'] = self.ReLU21
 236 | 
 237 |         self.conv22 = nn.Conv2d(
 238 |             in_channels = 128,
 239 |             out_channels = 128,
 240 |             kernel_size = 3,
 241 |             padding = 1
 242 |         )
 243 |         self.features.append(self.conv22)
 244 |         self.layerDict['conv22'] = self.conv22
 245 | 
 246 |         self.ReLU22 = nn.ReLU()
 247 |         self.features.append(self.ReLU22)
 248 |         self.layerDict['ReLU22'] = self.ReLU22
 249 | 
 250 |         self.pool2 = nn.MaxPool2d(2,2)
 251 |         self.features.append(self.pool2)
 252 |         self.layerDict['pool2'] = self.pool2
 253 | 
 254 |         self.conv31 = nn.Conv2d(
 255 |             in_channels = 128,
 256 |             out_channels = 128,
 257 |             kernel_size = 3,
 258 |             padding = 1
 259 |         )
 260 |         self.features.append(self.conv31)
 261 |         self.layerDict['conv31'] = self.conv31
 262 | 
 263 |         self.ReLU31 = nn.ReLU()
 264 |         self.features.append(self.ReLU31)
 265 |         self.layerDict['ReLU31'] = self.ReLU31
 266 | 
 267 |         self.conv32 = nn.Conv2d(
 268 |             in_channels = 128,
 269 |             out_channels = 128,
 270 |             kernel_size = 3,
 271 |             padding = 1
 272 |         )
 273 |         self.features.append(self.conv32)
 274 |         self.layerDict['conv32'] = self.conv32
 275 | 
 276 | 
 277 |         self.ReLU32 = nn.ReLU()
 278 |         self.features.append(self.ReLU32)
 279 |         self.layerDict['ReLU32'] = self.ReLU32
 280 | 
 281 |         self.pool3 = nn.MaxPool2d(2,2)
 282 |         self.features.append(self.pool3)
 283 |         self.layerDict['pool3'] = self.pool3
 284 | 
 285 |         self.classifier = []
 286 | 
 287 |         self.feature_dims = 4 * 4 * 128
 288 |         self.fc1 = nn.Linear(self.feature_dims, 512)
 289 |         self.classifier.append(self.fc1)
 290 |         self.layerDict['fc1'] = self.fc1
 291 | 
 292 |         self.fc1act = nn.Sigmoid()
 293 |         self.classifier.append(self.fc1act)
 294 |         self.layerDict['fc1act'] = self.fc1act
 295 | 
 296 |         self.fc2 = nn.Linear(512, 10)
 297 |         self.classifier.append(self.fc2)
 298 |         self.layerDict['fc2'] = self.fc2
 299 | 
 300 |     def forward(self, x):
 301 |         for layer in self.features:
 302 |             x = layer(x)
 303 | 
 304 |         x = x.view(-1, self.feature_dims)
 305 | 
 306 |         for layer in self.classifier:
 307 |             x = layer(x)
 308 |         return x
 309 | 
 310 | 
 311 |     def forward_from(self, x, layer):
 312 | 
 313 |         if layer in self.layerDict:
 314 |             targetLayer = self.layerDict[layer]
 315 | 
 316 |             if targetLayer in self.features:
 317 |                 layeridx = self.features.index(targetLayer)
 318 |                 for func in self.features[layeridx+1:]:
 319 |                     x = func(x)
 320 | 
 321 |                 x = x.view(-1, self.feature_dims)
 322 |                 for func in self.classifier:
 323 |                     x = func(x)
 324 |                 return x
 325 | 
 326 |             else:
 327 |                 layeridx = self.classifier.index(targetLayer)
 328 |                 for func in self.classifier[layeridx:]:
 329 |                     x = func(x)
 330 |                 return x
 331 |         else:
 332 |             print "layer not exists"
 333 |             exit(1)
 334 | 
 335 | 
 336 |     def getLayerOutput(self, x, targetLayer):
 337 |         for layer in self.features:
 338 |             x = layer(x)
 339 |             if layer == targetLayer:
 340 |                 return x
 341 | 
 342 |         x = x.view(-1, self.feature_dims)
 343 |         for layer in self.classifier:
 344 |             x = layer(x)
 345 |             if layer == targetLayer:
 346 |                 return x
 347 | 
 348 |         # Should not go here
 349 |         print "Target layer not found"
 350 |         exit(1)
 351 | 
 352 | 
 353 | class LeNetAlternativeconv1(nn.Module):
 354 |     def __init__(self, NChannels):
 355 |         super(LeNetAlternativeconv1, self).__init__()
 356 |         self.features = []
 357 |         self.classifier = []
 358 |         self.layerDict = collections.OrderedDict()
 359 | 
 360 |         self.conv0 = nn.Conv2d(
 361 |             in_channels = NChannels,
 362 |             out_channels = 8,
 363 |             kernel_size = 3
 364 |         )
 365 |         self.features.append(self.conv0)
 366 |         self.layerDict['conv0'] = self.conv0
 367 | 
 368 |         self.conv1 = nn.Conv2d(
 369 |             in_channels = 8,
 370 |             out_channels = 8,
 371 |             kernel_size = 3
 372 |         )
 373 |         self.features.append(self.conv1)
 374 |         self.layerDict['conv1'] = self.conv1
 375 | 
 376 |     def forward(self, x):
 377 |         for layer in self.features:
 378 |             x = layer(x)
 379 |         return x
 380 | 
 381 | 
 382 | class LeNetAlternativeconv1Archi2(nn.Module):
 383 |     def __init__(self, NChannels):
 384 |         super(LeNetAlternativeconv1Archi2, self).__init__()
 385 |         self.features = []
 386 |         self.classifier = []
 387 |         self.layerDict = collections.OrderedDict()
 388 | 
 389 |         self.conv0 = nn.Conv2d(
 390 |             in_channels = NChannels,
 391 |             out_channels = 8,
 392 |             kernel_size = 3
 393 |         )
 394 |         self.features.append(self.conv0)
 395 |         self.layerDict['conv0'] = self.conv0
 396 | 
 397 |         self.conv1 = nn.Conv2d(
 398 |             in_channels = 8,
 399 |             out_channels = 8,
 400 |             kernel_size = 3
 401 |         )
 402 |         self.features.append(self.conv1)
 403 |         self.layerDict['conv1'] = self.conv1
 404 | 
 405 |     def forward(self, x):
 406 |         for layer in self.features:
 407 |             x = layer(x)
 408 |         return x
 409 | 
 410 | 
 411 | class LeNetAlternativeReLU1(nn.Module):
 412 |     def __init__(self, NChannels):
 413 |         super(LeNetAlternativeReLU1, self).__init__()
 414 |         self.features = []
 415 |         self.classifier = []
 416 |         self.layerDict = collections.OrderedDict()
 417 | 
 418 |         self.conv1 = nn.Conv2d(
 419 |             in_channels = NChannels,
 420 |             out_channels = 8,
 421 |             kernel_size = 5
 422 |         )
 423 |         self.features.append(self.conv1)
 424 |         self.layerDict['conv1'] = self.conv1
 425 | 
 426 |         self.ReLU1 = nn.ReLU(True)
 427 |         self.features.append(self.ReLU1)
 428 |         self.layerDict['ReLU1'] = self.ReLU1
 429 | 
 430 |     def forward(self, x):
 431 |         for layer in self.features:
 432 |             x = layer(x)
 433 |         return x
 434 | 
 435 | 
 436 | class LeNetAlternativepool1(nn.Module):
 437 |     def __init__(self, NChannels):
 438 |         super(LeNetAlternativepool1, self).__init__()
 439 |         self.features = []
 440 |         self.classifier = []
 441 |         self.layerDict = collections.OrderedDict()
 442 | 
 443 |         self.conv1 = nn.Conv2d(
 444 |             in_channels = NChannels,
 445 |             out_channels = 8,
 446 |             kernel_size = 5
 447 |         )
 448 |         self.features.append(self.conv1)
 449 |         self.layerDict['conv1'] = self.conv1
 450 | 
 451 |         self.ReLU1 = nn.ReLU(True)
 452 |         self.features.append(self.ReLU1)
 453 |         self.layerDict['ReLU1'] = self.ReLU1
 454 | 
 455 |         self.pool1 = nn.MaxPool2d(2,2)
 456 |         self.features.append(self.pool1)
 457 |         self.layerDict['pool1'] = self.pool1
 458 | 
 459 |     def forward(self, x):
 460 |         for layer in self.features:
 461 |             x = layer(x)
 462 |         return x
 463 | 
 464 | class LeNetAlternativeconv2(nn.Module):
 465 |     def __init__(self, NChannels):
 466 |         super(LeNetAlternativeconv2, self).__init__()
 467 |         self.features = []
 468 |         self.classifier = []
 469 |         self.layerDict = collections.OrderedDict()
 470 | 
 471 |         self.conv1 = nn.Conv2d(
 472 |             in_channels = NChannels,
 473 |             out_channels = 8,
 474 |             kernel_size = 5
 475 |         )
 476 |         self.features.append(self.conv1)
 477 |         self.layerDict['conv1'] = self.conv1
 478 | 
 479 |         self.ReLU1 = nn.ReLU(True)
 480 |         self.features.append(self.ReLU1)
 481 |         self.layerDict['ReLU1'] = self.ReLU1
 482 | 
 483 |         self.pool1 = nn.MaxPool2d(2,2)
 484 |         self.features.append(self.pool1)
 485 |         self.layerDict['pool1'] = self.pool1
 486 | 
 487 |         self.conv2 = nn.Conv2d(
 488 |             in_channels = 8,
 489 |             out_channels = 16,
 490 |             kernel_size = 5
 491 |         )
 492 |         self.features.append(self.conv2)
 493 |         self.layerDict['conv2'] = self.conv2
 494 | 
 495 |     def forward(self, x):
 496 |         for layer in self.features:
 497 |             x = layer(x)
 498 |         return x
 499 | 
 500 | 
 501 | class LeNetAlternativeReLU2(nn.Module):
 502 |     def __init__(self, NChannels):
 503 |         super(LeNetAlternativeReLU2, self).__init__()
 504 |         self.features = []
 505 |         self.classifier = []
 506 |         self.layerDict = collections.OrderedDict()
 507 | 
 508 |         self.conv1 = nn.Conv2d(
 509 |             in_channels = NChannels,
 510 |             out_channels = 8,
 511 |             kernel_size = 5
 512 |         )
 513 |         self.features.append(self.conv1)
 514 |         self.layerDict['conv1'] = self.conv1
 515 | 
 516 |         self.ReLU1 = nn.ReLU(True)
 517 |         self.features.append(self.ReLU1)
 518 |         self.layerDict['ReLU1'] = self.ReLU1
 519 | 
 520 |         self.pool1 = nn.MaxPool2d(2,2)
 521 |         self.features.append(self.pool1)
 522 |         self.layerDict['pool1'] = self.pool1
 523 | 
 524 |         self.conv2 = nn.Conv2d(
 525 |             in_channels = 8,
 526 |             out_channels = 16,
 527 |             kernel_size = 5
 528 |         )
 529 |         self.features.append(self.conv2)
 530 |         self.layerDict['conv2'] = self.conv2
 531 | 
 532 |         self.ReLU2 = nn.ReLU(True)
 533 |         self.features.append(self.ReLU2)
 534 |         self.layerDict['ReLU2'] = self.ReLU2
 535 | 
 536 |     def forward(self, x):
 537 |         for layer in self.features:
 538 |             x = layer(x)
 539 |         return x
 540 | 
 541 | 
 542 | class LeNetAlternativepool2(nn.Module):
 543 |     def __init__(self, NChannels):
 544 |         super(LeNetAlternativepool2, self).__init__()
 545 |         self.features = []
 546 |         self.classifier = []
 547 |         self.layerDict = collections.OrderedDict()
 548 | 
 549 |         self.conv1 = nn.Conv2d(
 550 |             in_channels = NChannels,
 551 |             out_channels = 8,
 552 |             kernel_size = 5
 553 |         )
 554 |         self.features.append(self.conv1)
 555 |         self.layerDict['conv1'] = self.conv1
 556 | 
 557 |         self.ReLU1 = nn.ReLU(True)
 558 |         self.features.append(self.ReLU1)
 559 |         self.layerDict['ReLU1'] = self.ReLU1
 560 | 
 561 |         self.pool1 = nn.MaxPool2d(2,2)
 562 |         self.features.append(self.pool1)
 563 |         self.layerDict['pool1'] = self.pool1
 564 | 
 565 |         self.conv2 = nn.Conv2d(
 566 |             in_channels = 8,
 567 |             out_channels = 16,
 568 |             kernel_size = 5
 569 |         )
 570 |         self.features.append(self.conv2)
 571 |         self.layerDict['conv2'] = self.conv2
 572 | 
 573 |         self.ReLU2 = nn.ReLU(True)
 574 |         self.features.append(self.ReLU2)
 575 |         self.layerDict['ReLU2'] = self.ReLU2
 576 | 
 577 |         self.pool2 = nn.MaxPool2d(2,2)
 578 |         self.features.append(self.pool2)
 579 |         self.layerDict['pool2'] = self.pool2
 580 | 
 581 |     def forward(self, x):
 582 |         for layer in self.features:
 583 |             x = layer(x)
 584 |         return x
 585 | 
 586 | 
 587 | class LeNetAlternativefc1(nn.Module):
 588 |     def __init__(self, NChannels):
 589 |         super(LeNetAlternativefc1, self).__init__()
 590 |         self.features = []
 591 |         self.classifier = []
 592 |         self.layerDict = collections.OrderedDict()
 593 | 
 594 |         self.conv1 = nn.Conv2d(
 595 |             in_channels = NChannels,
 596 |             out_channels = 8,
 597 |             kernel_size = 5
 598 |         )
 599 |         self.features.append(self.conv1)
 600 |         self.layerDict['conv1'] = self.conv1
 601 | 
 602 |         self.ReLU1 = nn.ReLU(True)
 603 |         self.features.append(self.ReLU1)
 604 |         self.layerDict['ReLU1'] = self.ReLU1
 605 | 
 606 |         self.pool1 = nn.MaxPool2d(2,2)
 607 |         self.features.append(self.pool1)
 608 |         self.layerDict['pool1'] = self.pool1
 609 | 
 610 |         self.conv2 = nn.Conv2d(
 611 |             in_channels = 8,
 612 |             out_channels = 16,
 613 |             kernel_size = 5
 614 |         )
 615 |         self.features.append(self.conv2)
 616 |         self.layerDict['conv2'] = self.conv2
 617 | 
 618 |         self.ReLU2 = nn.ReLU(True)
 619 |         self.features.append(self.ReLU2)
 620 |         self.layerDict['ReLU2'] = self.ReLU2
 621 | 
 622 |         self.pool2 = nn.MaxPool2d(2,2)
 623 |         self.features.append(self.pool2)
 624 |         self.layerDict['pool2'] = self.pool2
 625 | 
 626 |         self.feature_dims = 16 * 4 * 4
 627 |         self.fc1 = nn.Linear(self.feature_dims, 120)
 628 |         self.classifier.append(self.fc1)
 629 |         self.layerDict['fc1'] = self.fc1
 630 | 
 631 | 
 632 |     def forward(self, x):
 633 |         for layer in self.features:
 634 |             x = layer(x)
 635 | 
 636 |         #print 'x.size', x.size()
 637 |         x = x.view(-1, self.feature_dims)
 638 | 
 639 |         for layer in self.classifier:
 640 |             x = layer(x)
 641 |         return x
 642 | 
 643 | class LeNetAlternativefc1act(nn.Module):
 644 |     def __init__(self, NChannels):
 645 |         super(LeNetAlternativefc1act, self).__init__()
 646 |         self.features = []
 647 |         self.classifier = []
 648 |         self.layerDict = collections.OrderedDict()
 649 | 
 650 |         self.conv1 = nn.Conv2d(
 651 |             in_channels = NChannels,
 652 |             out_channels = 8,
 653 |             kernel_size = 5
 654 |         )
 655 |         self.features.append(self.conv1)
 656 |         self.layerDict['conv1'] = self.conv1
 657 | 
 658 |         self.ReLU1 = nn.ReLU(True)
 659 |         self.features.append(self.ReLU1)
 660 |         self.layerDict['ReLU1'] = self.ReLU1
 661 | 
 662 |         self.pool1 = nn.MaxPool2d(2,2)
 663 |         self.features.append(self.pool1)
 664 |         self.layerDict['pool1'] = self.pool1
 665 | 
 666 |         self.conv2 = nn.Conv2d(
 667 |             in_channels = 8,
 668 |             out_channels = 16,
 669 |             kernel_size = 5
 670 |         )
 671 |         self.features.append(self.conv2)
 672 |         self.layerDict['conv2'] = self.conv2
 673 | 
 674 |         self.ReLU2 = nn.ReLU(True)
 675 |         self.features.append(self.ReLU2)
 676 |         self.layerDict['ReLU2'] = self.ReLU2
 677 | 
 678 |         self.pool2 = nn.MaxPool2d(2,2)
 679 |         self.features.append(self.pool2)
 680 |         self.layerDict['pool2'] = self.pool2
 681 | 
 682 |         self.feature_dims = 16 * 4 * 4
 683 |         self.fc1 = nn.Linear(self.feature_dims, 120)
 684 |         self.classifier.append(self.fc1)
 685 |         self.layerDict['fc1'] = self.fc1
 686 | 
 687 |         self.fc1act = nn.ReLU(True)
 688 |         self.classifier.append(self.fc1act)
 689 |         self.layerDict['fc1act'] = self.fc1act
 690 | 
 691 | 
 692 |     def forward(self, x):
 693 |         for layer in self.features:
 694 |             x = layer(x)
 695 | 
 696 |         #print 'x.size', x.size()
 697 |         x = x.view(-1, self.feature_dims)
 698 | 
 699 |         for layer in self.classifier:
 700 |             x = layer(x)
 701 |         return x
 702 | 
 703 | 
 704 | class LeNetAlternativefc2(nn.Module):
 705 |     def __init__(self, NChannels):
 706 |         super(LeNetAlternativefc2, self).__init__()
 707 |         self.features = []
 708 |         self.classifier = []
 709 |         self.layerDict = collections.OrderedDict()
 710 | 
 711 |         self.conv1 = nn.Conv2d(
 712 |             in_channels = NChannels,
 713 |             out_channels = 8,
 714 |             kernel_size = 5
 715 |         )
 716 |         self.features.append(self.conv1)
 717 |         self.layerDict['conv1'] = self.conv1
 718 | 
 719 |         self.ReLU1 = nn.ReLU(True)
 720 |         self.features.append(self.ReLU1)
 721 |         self.layerDict['ReLU1'] = self.ReLU1
 722 | 
 723 |         self.pool1 = nn.MaxPool2d(2,2)
 724 |         self.features.append(self.pool1)
 725 |         self.layerDict['pool1'] = self.pool1
 726 | 
 727 |         self.conv2 = nn.Conv2d(
 728 |             in_channels = 8,
 729 |             out_channels = 16,
 730 |             kernel_size = 5
 731 |         )
 732 |         self.features.append(self.conv2)
 733 |         self.layerDict['conv2'] = self.conv2
 734 | 
 735 |         self.ReLU2 = nn.ReLU(True)
 736 |         self.features.append(self.ReLU2)
 737 |         self.layerDict['ReLU2'] = self.ReLU2
 738 | 
 739 |         self.pool2 = nn.MaxPool2d(2,2)
 740 |         self.features.append(self.pool2)
 741 |         self.layerDict['pool2'] = self.pool2
 742 | 
 743 |         self.feature_dims = 16 * 4 * 4
 744 |         self.fc1 = nn.Linear(self.feature_dims, 120)
 745 |         self.classifier.append(self.fc1)
 746 |         self.layerDict['fc1'] = self.fc1
 747 | 
 748 |         self.fc1act = nn.ReLU(True)
 749 |         self.classifier.append(self.fc1act)
 750 |         self.layerDict['fc1act'] = self.fc1act
 751 | 
 752 | 
 753 |         self.fc2 = nn.Linear(120, 84)
 754 |         self.classifier.append(self.fc2)
 755 |         self.layerDict['fc2'] = self.fc2
 756 | 
 757 |     def forward(self, x):
 758 |         for layer in self.features:
 759 |             x = layer(x)
 760 | 
 761 |         #print 'x.size', x.size()
 762 |         x = x.view(-1, self.feature_dims)
 763 | 
 764 |         for layer in self.classifier:
 765 |             x = layer(x)
 766 |         return x
 767 | 
 768 | class LeNetAlternativefc2act(nn.Module):
 769 |     def __init__(self, NChannels):
 770 |         super(LeNetAlternativefc2act, self).__init__()
 771 |         self.features = []
 772 |         self.classifier = []
 773 |         self.layerDict = collections.OrderedDict()
 774 | 
 775 |         self.conv1 = nn.Conv2d(
 776 |             in_channels = NChannels,
 777 |             out_channels = 8,
 778 |             kernel_size = 5
 779 |         )
 780 |         self.features.append(self.conv1)
 781 |         self.layerDict['conv1'] = self.conv1
 782 | 
 783 |         self.ReLU1 = nn.ReLU(True)
 784 |         self.features.append(self.ReLU1)
 785 |         self.layerDict['ReLU1'] = self.ReLU1
 786 | 
 787 |         self.pool1 = nn.MaxPool2d(2,2)
 788 |         self.features.append(self.pool1)
 789 |         self.layerDict['pool1'] = self.pool1
 790 | 
 791 |         self.conv2 = nn.Conv2d(
 792 |             in_channels = 8,
 793 |             out_channels = 16,
 794 |             kernel_size = 5
 795 |         )
 796 |         self.features.append(self.conv2)
 797 |         self.layerDict['conv2'] = self.conv2
 798 | 
 799 |         self.ReLU2 = nn.ReLU(True)
 800 |         self.features.append(self.ReLU2)
 801 |         self.layerDict['ReLU2'] = self.ReLU2
 802 | 
 803 |         self.pool2 = nn.MaxPool2d(2,2)
 804 |         self.features.append(self.pool2)
 805 |         self.layerDict['pool2'] = self.pool2
 806 | 
 807 |         self.feature_dims = 16 * 4 * 4
 808 |         self.fc1 = nn.Linear(self.feature_dims, 120)
 809 |         self.classifier.append(self.fc1)
 810 |         self.layerDict['fc1'] = self.fc1
 811 | 
 812 |         self.fc1act = nn.ReLU(True)
 813 |         self.classifier.append(self.fc1act)
 814 |         self.layerDict['fc1act'] = self.fc1act
 815 | 
 816 | 
 817 |         self.fc2 = nn.Linear(120, 84)
 818 |         self.classifier.append(self.fc2)
 819 |         self.layerDict['fc2'] = self.fc2
 820 | 
 821 |         self.fc2act = nn.ReLU(True)
 822 |         self.classifier.append(self.fc2act)
 823 |         self.layerDict['fc2act'] = self.fc2act
 824 | 
 825 |     def forward(self, x):
 826 |         for layer in self.features:
 827 |             x = layer(x)
 828 | 
 829 |         #print 'x.size', x.size()
 830 |         x = x.view(-1, self.feature_dims)
 831 | 
 832 |         for layer in self.classifier:
 833 |             x = layer(x)
 834 |         return x
 835 | 
 836 | class LeNetAlternativefc3(nn.Module):
 837 |     def __init__(self, NChannels):
 838 |         super(LeNetAlternativefc3, self).__init__()
 839 |         self.features = []
 840 |         self.classifier = []
 841 |         self.layerDict = collections.OrderedDict()
 842 | 
 843 |         self.conv1 = nn.Conv2d(
 844 |             in_channels = NChannels,
 845 |             out_channels = 8,
 846 |             kernel_size = 5
 847 |         )
 848 |         self.features.append(self.conv1)
 849 |         self.layerDict['conv1'] = self.conv1
 850 | 
 851 |         self.ReLU1 = nn.ReLU(True)
 852 |         self.features.append(self.ReLU1)
 853 |         self.layerDict['ReLU1'] = self.ReLU1
 854 | 
 855 |         self.pool1 = nn.MaxPool2d(2,2)
 856 |         self.features.append(self.pool1)
 857 |         self.layerDict['pool1'] = self.pool1
 858 | 
 859 |         self.conv2 = nn.Conv2d(
 860 |             in_channels = 8,
 861 |             out_channels = 16,
 862 |             kernel_size = 5
 863 |         )
 864 |         self.features.append(self.conv2)
 865 |         self.layerDict['conv2'] = self.conv2
 866 | 
 867 |         self.ReLU2 = nn.ReLU(True)
 868 |         self.features.append(self.ReLU2)
 869 |         self.layerDict['ReLU2'] = self.ReLU2
 870 | 
 871 |         self.pool2 = nn.MaxPool2d(2,2)
 872 |         self.features.append(self.pool2)
 873 |         self.layerDict['pool2'] = self.pool2
 874 | 
 875 |         self.feature_dims = 16 * 4 * 4
 876 |         self.fc1 = nn.Linear(self.feature_dims, 120)
 877 |         self.classifier.append(self.fc1)
 878 |         self.layerDict['fc1'] = self.fc1
 879 | 
 880 |         self.fc1act = nn.ReLU(True)
 881 |         self.classifier.append(self.fc1act)
 882 |         self.layerDict['fc1act'] = self.fc1act
 883 | 
 884 | 
 885 |         self.fc2 = nn.Linear(120, 84)
 886 |         self.classifier.append(self.fc2)
 887 |         self.layerDict['fc2'] = self.fc2
 888 | 
 889 |         self.fc2act = nn.ReLU(True)
 890 |         self.classifier.append(self.fc2act)
 891 |         self.layerDict['fc2act'] = self.fc2act
 892 | 
 893 |         self.fc3 = nn.Linear(84, 10)
 894 |         self.classifier.append(self.fc3)
 895 |         self.layerDict['fc3'] = self.fc3
 896 | 
 897 | 
 898 |     def forward(self, x):
 899 |         for layer in self.features:
 900 |             x = layer(x)
 901 | 
 902 |         #print 'x.size', x.size()
 903 |         x = x.view(-1, self.feature_dims)
 904 | 
 905 |         for layer in self.classifier:
 906 |             x = layer(x)
 907 |         return x
 908 | 
 909 | 
 910 | 
 911 | class LeNetAlternativeReLU2Archi2(nn.Module):
 912 |     def __init__(self, NChannels):
 913 |         super(LeNetAlternativeReLU2Archi2, self).__init__()
 914 |         self.features = []
 915 |         self.classifier = []
 916 |         self.layerDict = collections.OrderedDict()
 917 | 
 918 |         self.conv11 = nn.Conv2d(
 919 |             in_channels = NChannels,
 920 |             out_channels = 8,
 921 |             kernel_size = 3
 922 |         )
 923 |         self.features.append(self.conv11)
 924 |         self.layerDict['conv11'] = self.conv11
 925 | 
 926 |         self.conv12 = nn.Conv2d(
 927 |             in_channels = 8,
 928 |             out_channels = 16,
 929 |             kernel_size = 3
 930 |         )
 931 |         self.features.append(self.conv12)
 932 |         self.layerDict['conv12'] = self.conv12
 933 | 
 934 |         self.pool1 = nn.MaxPool2d(2,2)
 935 |         self.features.append(self.pool1)
 936 |         self.layerDict['pool1'] = self.pool1
 937 | 
 938 |         self.conv21 = nn.Conv2d(
 939 |             in_channels = 16,
 940 |             out_channels = 16,
 941 |             kernel_size = 3
 942 |         )
 943 |         self.features.append(self.conv21)
 944 |         self.layerDict['conv21'] = self.conv21
 945 | 
 946 |         self.conv2 = nn.Conv2d(
 947 |             in_channels = 16,
 948 |             out_channels = 16,
 949 |             kernel_size = 3
 950 |         )
 951 |         self.features.append(self.conv2)
 952 |         self.layerDict['conv2'] = self.conv2
 953 | 
 954 |         self.ReLU2 = nn.ReLU(True)
 955 |         self.features.append(self.ReLU2)
 956 |         self.layerDict['ReLU2'] = self.ReLU2
 957 | 
 958 |     def forward(self, x):
 959 |         for layer in self.features:
 960 |             x = layer(x)
 961 |         return x
 962 | 
 963 | 
 964 | 
 965 | class CIFAR10CNNAlternativeconv11(nn.Module):
 966 |     def __init__(self, NChannels):
 967 |         super(CIFAR10CNNAlternativeconv11, self).__init__()
 968 |         self.features = []
 969 |         self.layerDict = collections.OrderedDict()
 970 | 
 971 |         self.conv11 = nn.Conv2d(
 972 |             in_channels = NChannels,
 973 |             out_channels = 64,
 974 |             kernel_size = 3,
 975 |             padding = 1
 976 |         )
 977 |         self.features.append(self.conv11)
 978 |         self.layerDict['conv11'] = self.conv11
 979 | 
 980 |     def forward(self, x):
 981 |         for layer in self.features:
 982 |             x = layer(x)
 983 |         return x
 984 | 
 985 |     def getLayerOutput(self, x, targetLayer):
 986 |         for layer in self.features:
 987 |             x = layer(x)
 988 |             if layer == targetLayer:
 989 |                 return x
 990 | 
 991 |         # Should not go here
 992 |         print "Target layer not found"
 993 |         exit(1)
 994 | 
 995 | 
 996 | class CIFAR10CNNAlternativeconv11Archi2(nn.Module):
 997 |     def __init__(self, NChannels):
 998 |         super(CIFAR10CNNAlternativeconv11Archi2, self).__init__()
 999 |         self.features = []
1000 |         self.layerDict = collections.OrderedDict()
1001 | 
1002 |         self.conv10 = nn.Conv2d(
1003 |             in_channels = NChannels,
1004 |             out_channels = 16,
1005 |             kernel_size = 3,
1006 |             padding = 1
1007 |         )
1008 |         self.features.append(self.conv10)
1009 |         self.layerDict['conv10'] = self.conv10
1010 | 
1011 |         self.conv11 = nn.Conv2d(
1012 |             in_channels = 16,
1013 |             out_channels = 64,
1014 |             kernel_size = 3,
1015 |             padding = 1
1016 |         )
1017 |         self.features.append(self.conv11)
1018 |         self.layerDict['conv11'] = self.conv11
1019 | 
1020 |     def forward(self, x):
1021 |         for layer in self.features:
1022 |             x = layer(x)
1023 |         return x
1024 | 
1025 |     def getLayerOutput(self, x, targetLayer):
1026 |         for layer in self.features:
1027 |             x = layer(x)
1028 |             if layer == targetLayer:
1029 |                 return x
1030 | 
1031 |         # Should not go here
1032 |         print "Target layer not found"
1033 |         exit(1)
1034 | 
1035 | 
1036 | class CIFAR10CNNAlternativeReLU22(nn.Module):
1037 |     def __init__(self, NChannels):
1038 |         super(CIFAR10CNNAlternativeReLU22, self).__init__()
1039 |         self.features = []
1040 |         self.layerDict = collections.OrderedDict()
1041 | 
1042 |         self.conv11 = nn.Conv2d(
1043 |             in_channels = NChannels,
1044 |             out_channels = 64,
1045 |             kernel_size = 3,
1046 |             padding = 1
1047 |         )
1048 |         self.features.append(self.conv11)
1049 |         self.layerDict['conv11'] = self.conv11
1050 | 
1051 |         self.ReLU11 = nn.ReLU()
1052 |         self.features.append(self.ReLU11)
1053 |         self.layerDict['ReLU11'] = self.ReLU11
1054 | 
1055 |         self.conv12 = nn.Conv2d(
1056 |             in_channels = 64,
1057 |             out_channels = 64,
1058 |             kernel_size = 3,
1059 |             padding = 1
1060 |         )
1061 |         self.features.append(self.conv12)
1062 |         self.layerDict['conv12'] = self.conv12
1063 | 
1064 |         self.ReLU12 = nn.ReLU()
1065 |         self.features.append(self.ReLU12)
1066 |         self.layerDict['ReLU12'] = self.ReLU12
1067 | 
1068 |         self.pool1 = nn.MaxPool2d(2,2)
1069 |         self.features.append(self.pool1)
1070 |         self.layerDict['pool1'] = self.pool1
1071 | 
1072 | 
1073 |         self.conv21 = nn.Conv2d(
1074 |             in_channels = 64,
1075 |             out_channels = 128,
1076 |             kernel_size = 3,
1077 |             padding = 1
1078 |         )
1079 |         self.features.append(self.conv21)
1080 |         self.layerDict['conv21'] = self.conv21
1081 | 
1082 |         self.ReLU21 = nn.ReLU()
1083 |         self.features.append(self.ReLU21)
1084 |         self.layerDict['ReLU21'] = self.ReLU21
1085 | 
1086 |         self.conv22 = nn.Conv2d(
1087 |             in_channels = 128,
1088 |             out_channels = 128,
1089 |             kernel_size = 3,
1090 |             padding = 1
1091 |         )
1092 |         self.features.append(self.conv22)
1093 |         self.layerDict['conv22'] = self.conv22
1094 | 
1095 |         self.ReLU22 = nn.ReLU()
1096 |         self.features.append(self.ReLU22)
1097 |         self.layerDict['ReLU22'] = self.ReLU22
1098 | 
1099 |     def forward(self, x):
1100 |         for layer in self.features:
1101 |             x = layer(x)
1102 |         return x
1103 | 
1104 |     def getLayerOutput(self, x, targetLayer):
1105 |         for layer in self.features:
1106 |             x = layer(x)
1107 |             if layer == targetLayer:
1108 |                 return x
1109 | 
1110 |         # Should not go here
1111 |         raise Exception("Target layer not found")
1112 | 
1113 | 
1114 | 
1115 | class CIFAR10CNNAlternativeReLU22Archi2(nn.Module):
1116 |     def __init__(self, NChannels):
1117 |         super(CIFAR10CNNAlternativeReLU22Archi2, self).__init__()
1118 |         self.features = []
1119 |         self.layerDict = collections.OrderedDict()
1120 | 
1121 |         self.conv11 = nn.Conv2d(
1122 |             in_channels = NChannels,
1123 |             out_channels = 64,
1124 |             kernel_size = 5,
1125 |             padding = 2
1126 |         )
1127 |         self.features.append(self.conv11)
1128 |         self.layerDict['conv11'] = self.conv11
1129 | 
1130 |         self.ReLU11 = nn.ReLU()
1131 |         self.features.append(self.ReLU11)
1132 |         self.layerDict['ReLU11'] = self.ReLU11
1133 | 
1134 |         self.pool1 = nn.MaxPool2d(2,2)
1135 |         self.features.append(self.pool1)
1136 |         self.layerDict['pool1'] = self.pool1
1137 | 
1138 |         self.conv22 = nn.Conv2d(
1139 |             in_channels = 64,
1140 |             out_channels = 128,
1141 |             kernel_size = 5,
1142 |             padding = 2
1143 |         )
1144 |         self.features.append(self.conv22)
1145 |         self.layerDict['conv22'] = self.conv22
1146 | 
1147 |         self.ReLU22 = nn.ReLU()
1148 |         self.features.append(self.ReLU22)
1149 |         self.layerDict['ReLU22'] = self.ReLU22
1150 | 
1151 |     def forward(self, x):
1152 |         for layer in self.features:
1153 |             x = layer(x)
1154 |         return x
1155 | 
1156 |     def getLayerOutput(self, x, targetLayer):
1157 |         for layer in self.features:
1158 |             x = layer(x)
1159 |             if layer == targetLayer:
1160 |                 return x
1161 | 
1162 |         # Should not go here
1163 |         raise Exception("Target layer not found")
1164 | 
1165 | 
1166 | class CIFAR10CNNAlternativeReLU32(nn.Module):
1167 |     def __init__(self, NChannels):
1168 |         super(CIFAR10CNNAlternativeReLU32, self).__init__()
1169 |         self.features = []
1170 |         self.layerDict = collections.OrderedDict()
1171 | 
1172 |         self.conv11 = nn.Conv2d(
1173 |             in_channels = NChannels,
1174 |             out_channels = 64,
1175 |             kernel_size = 3,
1176 |             padding = 1
1177 |         )
1178 |         self.features.append(self.conv11)
1179 |         self.layerDict['conv11'] = self.conv11
1180 | 
1181 |         self.ReLU11 = nn.ReLU()
1182 |         self.features.append(self.ReLU11)
1183 |         self.layerDict['ReLU11'] = self.ReLU11
1184 | 
1185 |         self.conv12 = nn.Conv2d(
1186 |             in_channels = 64,
1187 |             out_channels = 64,
1188 |             kernel_size = 3,
1189 |             padding = 1
1190 |         )
1191 |         self.features.append(self.conv12)
1192 |         self.layerDict['conv12'] = self.conv12
1193 | 
1194 |         self.ReLU12 = nn.ReLU()
1195 |         self.features.append(self.ReLU12)
1196 |         self.layerDict['ReLU12'] = self.ReLU12
1197 | 
1198 |         self.pool1 = nn.MaxPool2d(2,2)
1199 |         self.features.append(self.pool1)
1200 |         self.layerDict['pool1'] = self.pool1
1201 | 
1202 | 
1203 |         self.conv21 = nn.Conv2d(
1204 |             in_channels = 64,
1205 |             out_channels = 128,
1206 |             kernel_size = 3,
1207 |             padding = 1
1208 |         )
1209 |         self.features.append(self.conv21)
1210 |         self.layerDict['conv21'] = self.conv21
1211 | 
1212 |         self.ReLU21 = nn.ReLU()
1213 |         self.features.append(self.ReLU21)
1214 |         self.layerDict['ReLU21'] = self.ReLU21
1215 | 
1216 |         self.conv22 = nn.Conv2d(
1217 |             in_channels = 128,
1218 |             out_channels = 128,
1219 |             kernel_size = 3,
1220 |             padding = 1
1221 |         )
1222 |         self.features.append(self.conv22)
1223 |         self.layerDict['conv22'] = self.conv22
1224 | 
1225 |         self.ReLU22 = nn.ReLU()
1226 |         self.features.append(self.ReLU22)
1227 |         self.layerDict['ReLU22'] = self.ReLU22
1228 | 
1229 |         self.pool2 = nn.MaxPool2d(2,2)
1230 |         self.features.append(self.pool2)
1231 |         self.layerDict['pool2'] = self.pool2
1232 | 
1233 |         self.conv31 = nn.Conv2d(
1234 |             in_channels = 128,
1235 |             out_channels = 128,
1236 |             kernel_size = 3,
1237 |             padding = 1
1238 |         )
1239 |         self.features.append(self.conv31)
1240 |         self.layerDict['conv31'] = self.conv31
1241 | 
1242 |         self.ReLU31 = nn.ReLU()
1243 |         self.features.append(self.ReLU31)
1244 |         self.layerDict['ReLU31'] = self.ReLU31
1245 | 
1246 |         self.conv32 = nn.Conv2d(
1247 |             in_channels = 128,
1248 |             out_channels = 128,
1249 |             kernel_size = 3,
1250 |             padding = 1
1251 |         )
1252 |         self.features.append(self.conv32)
1253 |         self.layerDict['conv32'] = self.conv32
1254 | 
1255 | 
1256 |         self.ReLU32 = nn.ReLU()
1257 |         self.features.append(self.ReLU32)
1258 |         self.layerDict['ReLU32'] = self.ReLU32
1259 | 
1260 |     def forward(self, x):
1261 |         for layer in self.features:
1262 |             x = layer(x)
1263 |         return x
1264 | 
1265 |     def getLayerOutput(self, x, targetLayer):
1266 |         for layer in self.features:
1267 |             x = layer(x)
1268 |             if layer == targetLayer:
1269 |                 return x
1270 | 
1271 |         # Should not go here
1272 |         raise Exception("Target layer not found")
1273 | 
1274 | 
1275 | class CIFAR10CNNAlternativeReLU32Archi2(nn.Module):
1276 |     def __init__(self, NChannels):
1277 |         super(CIFAR10CNNAlternativeReLU32Archi2, self).__init__()
1278 |         self.features = []
1279 |         self.layerDict = collections.OrderedDict()
1280 | 
1281 |         self.conv11 = nn.Conv2d(
1282 |             in_channels = NChannels,
1283 |             out_channels = 64,
1284 |             kernel_size = 5,
1285 |             padding = 2
1286 |         )
1287 |         self.features.append(self.conv11)
1288 |         self.layerDict['conv11'] = self.conv11
1289 | 
1290 |         self.ReLU11 = nn.ReLU()
1291 |         self.features.append(self.ReLU11)
1292 |         self.layerDict['ReLU11'] = self.ReLU11
1293 | 
1294 |         self.pool1 = nn.MaxPool2d(2,2)
1295 |         self.features.append(self.pool1)
1296 |         self.layerDict['pool1'] = self.pool1
1297 | 
1298 |         self.conv22 = nn.Conv2d(
1299 |             in_channels = 64,
1300 |             out_channels = 128,
1301 |             kernel_size = 5,
1302 |             padding = 2
1303 |         )
1304 |         self.features.append(self.conv22)
1305 |         self.layerDict['conv22'] = self.conv22
1306 | 
1307 |         self.ReLU22 = nn.ReLU()
1308 |         self.features.append(self.ReLU22)
1309 |         self.layerDict['ReLU22'] = self.ReLU22
1310 | 
1311 |         self.pool2 = nn.MaxPool2d(2,2)
1312 |         self.features.append(self.pool2)
1313 |         self.layerDict['pool2'] = self.pool2
1314 | 
1315 |         self.conv32 = nn.Conv2d(
1316 |             in_channels = 128,
1317 |             out_channels = 128,
1318 |             kernel_size = 5,
1319 |             padding = 2
1320 |         )
1321 |         self.features.append(self.conv32)
1322 |         self.layerDict['conv32'] = self.conv32
1323 | 
1324 |         self.ReLU32 = nn.ReLU()
1325 |         self.features.append(self.ReLU32)
1326 |         self.layerDict['ReLU32'] = self.ReLU32
1327 | 
1328 |     def forward(self, x):
1329 |         for layer in self.features:
1330 |             x = layer(x)
1331 |         return x
1332 | 
1333 |     def getLayerOutput(self, x, targetLayer):
1334 |         for layer in self.features:
1335 |             x = layer(x)
1336 |             if layer == targetLayer:
1337 |                 return x
1338 | 
1339 |         # Should not go here
1340 |         raise Exception("Target layer not found")
1341 | 
1342 | 
1343 | 
1344 | class MLPAlternative(nn.Module):
1345 |     def __init__(self, inputSize, outputSize):
1346 |         super(MLPAlternative, self).__init__()
1347 |         self.features = []
1348 |         self.classifier = []
1349 |         self.layerDict = collections.OrderedDict()
1350 | 
1351 |         assert len(inputSize) == 3
1352 |         assert len(outputSize) == 3
1353 | 
1354 |         self.inputSize = inputSize
1355 |         self.outputSize = outputSize
1356 | 
1357 |         self.input_dims = 1
1358 |         for d in inputSize:
1359 |             self.input_dims *= d
1360 | 
1361 |         self.output_dims = 1
1362 |         for d in outputSize:
1363 |             self.output_dims *= d
1364 | 
1365 |         self.fc1_N = int(0.7*self.input_dims)
1366 |         self.fc2_N = int(0.1*self.input_dims)
1367 |         self.fc3_N = self.output_dims
1368 | 
1369 |         print "fc1_N: ", self.fc1_N
1370 |         print "fc2_N: ", self.fc2_N
1371 |         print "fc3_N: ", self.fc3_N
1372 | 
1373 |         self.fc1 = nn.Linear(self.input_dims, self.fc1_N)
1374 |         self.classifier.append(self.fc1)
1375 |         self.layerDict['fc1'] = self.fc1
1376 |         self.fc1act = nn.ReLU()
1377 |         self.classifier.append(self.fc1act)
1378 |         self.layerDict['fc1act'] = self.fc1act
1379 | 
1380 |         self.fc2 = nn.Linear(self.fc1_N, self.fc2_N)
1381 |         self.classifier.append(self.fc2)
1382 |         self.layerDict['fc2'] = self.fc2
1383 |         self.fc2act = nn.ReLU()
1384 |         self.classifier.append(self.fc2act)
1385 |         self.layerDict['fc2act'] = self.fc2act
1386 | 
1387 | 
1388 |         self.fc3 = nn.Linear(self.fc2_N, self.fc3_N)
1389 |         self.classifier.append(self.fc3)
1390 |         self.layerDict['fc3'] = self.fc3
1391 |         self.fc3act = nn.ReLU()
1392 |         self.classifier.append(self.fc3act)
1393 |         self.layerDict['fc3act'] = self.fc3act
1394 | 
1395 | 
1396 |     def forward(self, x):
1397 |         x = x.view(-1, self.input_dims)
1398 |         for layer in self.classifier:
1399 |             x = layer(x)
1400 |         x = x.view(-1, self.outputSize[0], self.outputSize[1], self.outputSize[2])
1401 |         return x
1402 | 
1403 |     def getLayerOutput(self, x, targetLayer):
1404 |         x = x.view(-1, self.input_dims)
1405 |         for layer in self.classifier:
1406 |             x = layer(x)
1407 |             if layer == targetLayer:
1408 |                 return x
1409 | 
1410 |         # Should not go here
1411 |         raise Exception("Target layer not found")
1412 | 
1413 | 
1414 | class LeNetDecoderconv1(nn.Module):
1415 |     def __init__(self, NChannels):
1416 |         super(LeNetDecoderconv1, self).__init__()
1417 |         self.decoder = []
1418 |         self.layerDict = collections.OrderedDict()
1419 | 
1420 |         self.conv1 = nn.Conv2d(
1421 |             in_channels = NChannels,
1422 |             out_channels = 1,
1423 |             kernel_size = 5,
1424 |             padding = (4,4)
1425 |         )
1426 |         self.layerDict['conv1'] = self.conv1
1427 | 
1428 |     def forward(self, x):
1429 |         for layer in self.layerDict:
1430 |             x = self.layerDict[layer](x)
1431 |         return x
1432 | 
1433 | 
1434 | 
1435 | class LeNetDecoderReLU2(nn.Module):
1436 |     def __init__(self, NChannels):
1437 |         super(LeNetDecoderReLU2, self).__init__()
1438 |         self.decoder = []
1439 |         self.layerDict = collections.OrderedDict()
1440 | 
1441 | 
1442 |         self.deconv1 = nn.ConvTranspose2d(
1443 |             in_channels = 16,
1444 |             out_channels = 8,
1445 |             kernel_size = 5,
1446 |         )
1447 | 
1448 |         self.layerDict['deconv1'] = self.deconv1
1449 | 
1450 |         self.ReLU1 = nn.ReLU()
1451 |         self.layerDict['ReLU1'] = self.ReLU1
1452 | 
1453 |         self.deconv2 = nn.ConvTranspose2d(
1454 |             in_channels = 8,
1455 |             out_channels = 1,
1456 |             kernel_size = 5,
1457 |             stride = 2,
1458 |             output_padding = (1,1)
1459 |         )
1460 |         self.layerDict['deconv2'] = self.deconv2
1461 | 
1462 | 
1463 |     def forward(self, x):
1464 |         for layer in self.layerDict:
1465 |             x = self.layerDict[layer](x)
1466 |         return x
1467 | 
1468 | 
1469 | 
1470 | class CIFAR10CNNDecoderconv1(nn.Module):
1471 |     def __init__(self, NChannels):
1472 |         super(CIFAR10CNNDecoderconv1, self).__init__()
1473 |         self.decoder = []
1474 |         self.layerDict = collections.OrderedDict()
1475 |         self.deconv1 = nn.ConvTranspose2d(
1476 |             in_channels = 16,
1477 |             out_channels = 8,
1478 |             kernel_size = 5,
1479 |         )
1480 | 
1481 |         self.layerDict['deconv1'] = self.deconv1
1482 |         self.ReLU1 = nn.ReLU()
1483 |         self.layerDict['ReLU1'] = self.ReLU1
1484 |         self.deconv2 = nn.ConvTranspose2d(
1485 |             in_channels = 8,
1486 |             out_channels = 1,
1487 |             kernel_size = 5,
1488 |             stride = 2,
1489 |             output_padding = (1,1)
1490 |         )
1491 |         self.layerDict['deconv2'] = self.deconv2
1492 |     def forward(self, x):
1493 |         for layer in self.layerDict:
1494 |             x = self.layerDict[layer](x)
1495 |         return x
1496 | 
1497 | 
1498 | 
1499 | class CIFAR10CNNDecoderconv1(nn.Module):
1500 |     def __init__(self, NChannels):
1501 |         super(CIFAR10CNNDecoderconv1, self).__init__()
1502 |         self.decoder = []
1503 |         self.layerDict = collections.OrderedDict()
1504 | 
1505 |         self.deconv11 = nn.ConvTranspose2d(
1506 |             in_channels = 64,
1507 |             out_channels = 3,
1508 |             kernel_size = 3,
1509 |             padding = 1
1510 |         )
1511 | 
1512 |         self.layerDict['deconv11'] = self.deconv11
1513 | 
1514 | 
1515 |     def forward(self, x):
1516 |         for layer in self.layerDict:
1517 |             x = self.layerDict[layer](x)
1518 |         return x
1519 | 
1520 | 
1521 | class CIFAR10CNNDecoderReLU22(nn.Module):
1522 |     def __init__(self, NChannels):
1523 |         super(CIFAR10CNNDecoderReLU22, self).__init__()
1524 | 
1525 |         self.layerDict = collections.OrderedDict()
1526 | 
1527 |         self.deconv11 = nn.ConvTranspose2d(
1528 |             in_channels = 128,
1529 |             out_channels = 64,
1530 |             kernel_size = 3,
1531 |             stride = 2,
1532 |             padding = 1,
1533 |             output_padding = 1
1534 |         )
1535 |         self.layerDict['deconv11'] = self.deconv11
1536 | 
1537 |         self.ReLU11 = nn.ReLU()
1538 |         self.layerDict['ReLU11'] = self.ReLU11
1539 | 
1540 |         self.deconv21 = nn.ConvTranspose2d(
1541 |             in_channels = 64,
1542 |             out_channels = 3,
1543 |             kernel_size = 3,
1544 |             padding = 1
1545 |         )
1546 | 
1547 |         self.layerDict['deconv21'] = self.deconv21
1548 | 
1549 |     def forward(self, x):
1550 |         for layer in self.layerDict:
1551 |             x = self.layerDict[layer](x)
1552 |         return x
1553 | 
1554 | 
1555 | class CIFAR10CNNDecoderReLU32(nn.Module):
1556 |     def __init__(self, NChannels):
1557 |         super(CIFAR10CNNDecoderReLU32, self).__init__()
1558 | 
1559 |         self.layerDict = collections.OrderedDict()
1560 | 
1561 |         self.deconv11 = nn.ConvTranspose2d(
1562 |             in_channels = 128,
1563 |             out_channels = 128,
1564 |             kernel_size = 3,
1565 |             stride = 2,
1566 |             padding = 1,
1567 |             output_padding = 1
1568 |         )
1569 |         self.layerDict['deconv11'] = self.deconv11
1570 | 
1571 |         self.ReLU11 = nn.ReLU()
1572 |         self.layerDict['ReLU11'] = self.ReLU11
1573 | 
1574 |         self.deconv21 = nn.ConvTranspose2d(
1575 |             in_channels = 128,
1576 |             out_channels = 64,
1577 |             kernel_size = 3,
1578 |             stride = 2,
1579 |             padding = 1,
1580 |             output_padding = 1
1581 |         )
1582 |         self.layerDict['deconv21'] = self.deconv21
1583 | 
1584 |         self.ReLU21 = nn.ReLU()
1585 |         self.layerDict['ReLU21'] = self.ReLU21
1586 | 
1587 |         self.deconv31 = nn.ConvTranspose2d(
1588 |             in_channels = 64,
1589 |             out_channels = 3,
1590 |             kernel_size = 3,
1591 |             padding = 1
1592 |         )
1593 | 
1594 |         self.layerDict['deconv31'] = self.deconv31
1595 | 
1596 |     def forward(self, x):
1597 |         for layer in self.layerDict:
1598 |             x = self.layerDict[layer](x)
1599 |         return x
1600 | 


--------------------------------------------------------------------------------
/net.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zechenghe/Inverse_Collaborative_Inference/f010aed0035238c573037b04c47c334ea733d15c/net.pyc


--------------------------------------------------------------------------------
/noise_generation.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2020-04-20
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | from net import *
 17 | from utils import *
 18 | 
 19 | from skimage.measure import compare_ssim
 20 | 
 21 | #####################
 22 | 
 23 | # This function is used to generate noise that does not affect model output
 24 | 
 25 | #####################
 26 | 
 27 | def noise_gen(args, model_dir = "checkpoints/MNIST/", model_name = "ckpt.pth"):
 28 | 
 29 |     sourceLayer = args.noise_sourceLayer
 30 |     targetLayer = args.noise_targetLayer
 31 |     gpu = args.gpu
 32 | 
 33 |     if args.dataset == 'MNIST':
 34 | 
 35 |         mu = torch.tensor([0.5], dtype=torch.float32)
 36 |         sigma = torch.tensor([0.5], dtype=torch.float32)
 37 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 38 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 39 | 
 40 |         tsf = {
 41 |             'train': transforms.Compose(
 42 |             [
 43 |                 transforms.ToTensor(),
 44 |                 Normalize
 45 |             ]),
 46 | 
 47 |             'test': transforms.Compose(
 48 |             [
 49 |                 transforms.ToTensor(),
 50 |                 Normalize
 51 |             ])
 52 |         }
 53 | 
 54 |         trainset = torchvision.datasets.MNIST(root='./data/MNIST', train=True,
 55 |                                         download=True, transform = tsf['train'])
 56 | 
 57 |         testset = torchvision.datasets.MNIST(root='./data/MNIST', train=False,
 58 |                                        download=True, transform = tsf['test'])
 59 | 
 60 |     x_train, y_train = trainset.data, trainset.targets,
 61 |     x_test, y_test = testset.data, testset.targets,
 62 | 
 63 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = 1,
 64 |                                       shuffle = False, num_workers = 1)
 65 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 66 |                                       shuffle = False, num_workers = 1)
 67 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
 68 |                                       shuffle = False, num_workers = 1)
 69 |     trainIter = iter(trainloader)
 70 |     testIter = iter(testloader)
 71 |     inverseIter = iter(inverseloader)
 72 | 
 73 |     net = torch.load(model_dir + model_name)
 74 |     if not gpu:
 75 |         net = net.cpu()
 76 | 
 77 |     net.eval()
 78 | 
 79 |     # Only to get the feature size
 80 |     targetImg, _ = getImgByClass(inverseIter, C = 0)
 81 |     deprocessImg = deprocess(targetImg.clone())
 82 | 
 83 |     softmaxLayer = nn.Softmax().cuda() if gpu else nn.Softmax()
 84 |     ReLULayer = nn.ReLU(False).cuda() if gpu else nn.ReLU(False)
 85 |     if gpu:
 86 |         targetImg = targetImg.cuda()
 87 | 
 88 |     layer = net.layerDict[sourceLayer]
 89 |     sourceLayerOutput = net.getLayerOutput(targetImg, layer)
 90 |     xGen = torch.ones(sourceLayerOutput.size(), requires_grad = True, device="cuda" if args.gpu else 'cpu')
 91 | 
 92 |     refSource = torch.randn(size=xGen.size(), requires_grad = True) * args.noise_level
 93 | 
 94 |     # If noise for relu layer, make all entries non-negtive
 95 |     if 'ReLU' in args.noise_sourceLayer:
 96 |         refSource = ReLULayer(refSource)
 97 | 
 98 | 
 99 |     layer = net.layerDict[targetLayer]
100 |     targetLayerOutput = net.getLayerOutput(targetImg, layer)
101 |     refTarget = torch.zeros(targetLayerOutput.size(), requires_grad = True)
102 | 
103 |     if args.gpu:
104 |         refTarget = refTarget.cuda()
105 |         refSource = refSource.cuda()
106 | 
107 |     #print "xGen.size", xGen.size()
108 |     #print "refSource.size", refSource.size()
109 |     #print "refTarget.size", refTarget.size()
110 | 
111 |     #targetLayerOutput = net.getLayerOutputFrom(
112 |     #    x = xGen,
113 |     #    sourceLayer = sourceLayer,
114 |     #    targetLayer = targetLayer
115 |     #)
116 | 
117 |     #print "targetLayerOutput.size", targetLayerOutput.size()
118 | 
119 |     optimizer = optim.Adam(
120 |         params = [xGen],
121 |         lr = args.noise_learning_rate,
122 |         eps = args.noise_eps,
123 |         amsgrad = args.noise_AMSGrad
124 |     )
125 | 
126 |     for i in range(args.noise_iters):
127 | 
128 |         optimizer.zero_grad()
129 | 
130 |         targetLayerOutput = net.getLayerOutputFrom(
131 |             x = ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen,
132 |             sourceLayer = sourceLayer,
133 |             targetLayer = targetLayer
134 |         )
135 | 
136 |         sourceLayerLoss = (((ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen) - refSource)**2).mean()
137 |         #sourceLayerLoss = -((ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen)**2).mean()
138 |         #sourceLayerLoss = -torch.abs(ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen).mean()
139 | 
140 | 
141 |         targetLayerLoss = ((targetLayerOutput - refTarget)**2).mean()
142 | 
143 |         totalLoss = targetLayerLoss + sourceLayerLoss * args.noise_lambda_sourcelayer
144 | 
145 |         totalLoss.backward(retain_graph=True)
146 |         optimizer.step()
147 | 
148 |         #print "Iter ", i, "loss: ", totalLoss.cpu().detach().numpy(), \
149 |         #"sourceLayerLoss: ", sourceLayerLoss.cpu().detach().numpy(), \
150 |         #"targetLayerLoss: ", targetLayerLoss.cpu().detach().numpy()
151 | 
152 |     noise_gen = (ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen).detach().cpu().numpy()
153 |     noise_dir = 'noise/' + args.dataset + '/'
154 |     noise_file_name = args.noise_sourceLayer + '-' + args.noise_targetLayer + '-' + str(round(args.noise_level, 2))
155 | 
156 |     #print noise_gen
157 |     #print sum(noise_gen)
158 | 
159 |     if not os.path.exists(noise_dir):
160 |         os.makedirs(noise_dir)
161 |     np.save(noise_dir + noise_file_name, noise_gen)
162 | 
163 |     acc = evalTestSplitModel(
164 |             testloader, net, net,
165 |             layer=args.noise_sourceLayer,
166 |             gpu = args.gpu,
167 |             noise_type = 'noise_gen',
168 |             noise_level = args.noise_level,
169 |             args = args
170 |         )
171 |     print "noise level", args.noise_level, "acc", acc
172 | 
173 |     return noise_gen
174 | 
175 | 
176 | if __name__ == '__main__':
177 |     import argparse
178 |     import sys
179 |     import traceback
180 | 
181 |     try:
182 |         parser = argparse.ArgumentParser()
183 |         parser.add_argument('--dataset', type = str, default = 'MNIST')
184 |         parser.add_argument('--network', type = str, default = 'LeNet')
185 |         parser.add_argument('--inverseClass', type = int, default = 0)
186 | 
187 |         parser.add_argument('--noise_iters', type = int, default = 500)
188 |         parser.add_argument('--noise_eps', type = float, default = 1e-3)
189 |         parser.add_argument('--noise_AMSGrad', type = bool, default = True)
190 |         parser.add_argument('--noise_learning_rate', type = float, default = 1e-1)
191 |         parser.add_argument('--noise_lambda_sourcelayer', type = float, default = 1e-1)
192 |         parser.add_argument('--noise_decrease_LR', type = int, default = 20)
193 |         parser.add_argument('--noise_sourceLayer', type = str, default = 'ReLU2')
194 |         parser.add_argument('--noise_targetLayer', type = str, default = 'fc3')
195 |         parser.add_argument('--noise_level', type = float, default = None)
196 | 
197 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
198 |         parser.set_defaults(gpu=True)
199 | 
200 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
201 |         parser.set_defaults(validation=True)
202 |         args = parser.parse_args()
203 | 
204 |         args.model_dir = "checkpoints/" + args.dataset + '/'
205 |         args.model_name = "ckpt.pth"
206 | 
207 |         if args.noise_level == None:
208 |             for nl in np.arange(0, 5, 0.5):
209 |                 args.noise_level = nl
210 |                 noise_gen(
211 |                     args = args,
212 |                     model_dir = args.model_dir,
213 |                     model_name = args.model_name
214 |                 )
215 |         else:
216 |             noise_gen(
217 |                 args = args,
218 |                 model_dir = args.model_dir,
219 |                 model_name = args.model_name
220 |             )
221 | 
222 |     except:
223 |         traceback.print_exc(file=sys.stdout)
224 |         sys.exit(1)
225 | 


--------------------------------------------------------------------------------
/noise_generation_opt.py:
--------------------------------------------------------------------------------
  1 | # @Author: Zecheng He
  2 | # @Date:   2020-04-20
  3 | 
  4 | import time
  5 | import math
  6 | import os
  7 | import numpy as np
  8 | 
  9 | import torch
 10 | import torchvision
 11 | import torchvision.transforms as transforms
 12 | import torch.nn as nn
 13 | import torch.nn.functional as F
 14 | import torch.optim as optim
 15 | import torch.backends.cudnn as cudnn
 16 | from net import *
 17 | from utils import *
 18 | 
 19 | from skimage.measure import compare_ssim
 20 | 
 21 | #####################
 22 | 
 23 | # This function is used to generate noise that does not affect model output
 24 | 
 25 | #####################
 26 | 
 27 | def noise_gen(args, model_dir = "checkpoints/MNIST/", model_name = "ckpt.pth"):
 28 | 
 29 |     sourceLayer = args.noise_sourceLayer
 30 |     targetLayer = args.noise_targetLayer
 31 |     gpu = args.gpu
 32 | 
 33 |     if args.dataset == 'MNIST':
 34 | 
 35 |         mu = torch.tensor([0.5], dtype=torch.float32)
 36 |         sigma = torch.tensor([0.5], dtype=torch.float32)
 37 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 38 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 39 | 
 40 |         tsf = {
 41 |             'train': transforms.Compose(
 42 |             [
 43 |                 transforms.ToTensor(),
 44 |                 Normalize
 45 |             ]),
 46 | 
 47 |             'test': transforms.Compose(
 48 |             [
 49 |                 transforms.ToTensor(),
 50 |                 Normalize
 51 |             ])
 52 |         }
 53 | 
 54 |         trainset = torchvision.datasets.MNIST(root='./data/MNIST', train=True,
 55 |                                         download=True, transform = tsf['train'])
 56 | 
 57 |         testset = torchvision.datasets.MNIST(root='./data/MNIST', train=False,
 58 |                                        download=True, transform = tsf['test'])
 59 | 
 60 |     x_train, y_train = trainset.data, trainset.targets,
 61 |     x_test, y_test = testset.data, testset.targets,
 62 | 
 63 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = args.noise_batch_size,
 64 |                                       shuffle = False, num_workers = 1)
 65 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
 66 |                                       shuffle = False, num_workers = 1)
 67 |     inverseloader = torch.utils.data.DataLoader(testset, batch_size = 1,
 68 |                                       shuffle = False, num_workers = 1)
 69 |     trainIter = iter(trainloader)
 70 |     testIter = iter(testloader)
 71 |     inverseIter = iter(inverseloader)
 72 | 
 73 |     net = torch.load(model_dir + model_name)
 74 |     if not gpu:
 75 |         net = net.cpu()
 76 | 
 77 |     net.eval()
 78 | 
 79 |     # Only to get the feature size
 80 |     targetImg, _ = getImgByClass(inverseIter, C = 0)
 81 |     deprocessImg = deprocess(targetImg.clone())
 82 | 
 83 |     softmaxLayer = nn.Softmax().cuda() if gpu else nn.Softmax()
 84 |     ReLULayer = nn.ReLU(False).cuda() if gpu else nn.ReLU(False)
 85 |     if gpu:
 86 |         targetImg = targetImg.cuda()
 87 | 
 88 |     layer = net.layerDict[sourceLayer]
 89 |     sourceLayerOutput = net.getLayerOutput(targetImg, layer)
 90 |     xGen = torch.ones(sourceLayerOutput.size(), requires_grad = True, device="cuda" if args.gpu else 'cpu')
 91 | 
 92 |     refSource = torch.randn(size=xGen.size(), requires_grad = True) * args.noise_level
 93 | 
 94 |     # If noise for relu layer, make all entries non-negtive
 95 |     #if 'ReLU' in args.noise_sourceLayer:
 96 |     #    refSource = ReLULayer(refSource)
 97 | 
 98 | 
 99 |     layer = net.layerDict[targetLayer]
100 |     targetLayerOutput = net.getLayerOutput(targetImg, layer)
101 | 
102 |     if args.gpu:
103 |         refSource = refSource.cuda()
104 | 
105 |     optimizer = optim.Adam(
106 |         params = [xGen],
107 |         lr = args.noise_learning_rate,
108 |         eps = args.noise_eps,
109 |         amsgrad = args.noise_AMSGrad
110 |     )
111 | 
112 | 
113 |     NBatch = len(trainset) / args.noise_batch_size
114 |     cudnn.benchmark = True
115 | 
116 | 
117 |     for epoch in range(args.noise_epochs):
118 |         lossTrain = 0.0
119 |         accTrain = 0.0
120 |         for i in range(NBatch):
121 |             try:
122 |                 batchX, batchY = trainIter.next()
123 |             except StopIteration:
124 |                 trainIter = iter(trainloader)
125 |                 batchX, batchY = trainIter.next()
126 | 
127 |             if gpu:
128 |                 batchX = batchX.cuda()
129 |                 batchY = batchY.cuda()
130 | 
131 |             optimizer.zero_grad()
132 | 
133 |             sourceLayerOutput = net.getLayerOutput(
134 |                 x = batchX,
135 |                 targetLayer = net.layerDict[sourceLayer]
136 |             )
137 | 
138 |             #print "sourceLayerOutput.size", sourceLayerOutput.size()
139 |             #print "xGen.size", xGen.size()
140 | 
141 |             targetLayerOutput = net.getLayerOutputFrom(
142 |                 x = sourceLayerOutput + torch.cat(args.noise_batch_size * [xGen]),
143 |                 sourceLayer = sourceLayer,
144 |                 targetLayer = targetLayer
145 |             )
146 | 
147 |             refTarget = net.getLayerOutput(
148 |                 x = batchX,
149 |                 targetLayer = net.layerDict[targetLayer]
150 |             )
151 | 
152 |             sourceLayerLoss = ((xGen - refSource)**2).mean()
153 |             #sourceLayerLoss = -((ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen)**2).mean()
154 |             #sourceLayerLoss = -torch.abs(ReLULayer(xGen) if 'ReLU' in args.noise_sourceLayer else xGen).mean()
155 | 
156 |             targetLayerLoss = ((targetLayerOutput - refTarget)**2).mean()
157 | 
158 |             totalLoss = targetLayerLoss + sourceLayerLoss * args.noise_lambda_sourcelayer
159 | 
160 |             totalLoss.backward(retain_graph=True)
161 |             optimizer.step()
162 | 
163 |         print "Epoch", epoch, "loss: ", totalLoss.cpu().detach().numpy(), \
164 |             "sourceLayerLoss: ", sourceLayerLoss.cpu().detach().numpy(), \
165 |             "targetLayerLoss: ", targetLayerLoss.cpu().detach().numpy()
166 | 
167 |     noise_gen = xGen.detach().cpu().numpy()
168 |     noise_dir = 'noise_opt/' + args.dataset + '/'
169 |     noise_file_name = args.noise_sourceLayer + '-' + args.noise_targetLayer + '-' + str(round(args.noise_level, 2))
170 | 
171 |     #print noise_gen
172 |     #print sum(noise_gen)
173 | 
174 |     if not os.path.exists(noise_dir):
175 |         os.makedirs(noise_dir)
176 |     np.save(noise_dir + noise_file_name, noise_gen)
177 | 
178 |     acc = evalTestSplitModel(
179 |             testloader, net, net,
180 |             layer=args.noise_sourceLayer,
181 |             gpu = args.gpu,
182 |             noise_type = 'noise_gen_opt',
183 |             noise_level = args.noise_level,
184 |             args = args
185 |         )
186 |     print "noise level", args.noise_level, "acc", acc
187 | 
188 |     return noise_gen
189 | 
190 | 
191 | if __name__ == '__main__':
192 |     import argparse
193 |     import sys
194 |     import traceback
195 | 
196 |     try:
197 |         parser = argparse.ArgumentParser()
198 |         parser.add_argument('--dataset', type = str, default = 'MNIST')
199 |         parser.add_argument('--network', type = str, default = 'LeNet')
200 |         parser.add_argument('--inverseClass', type = int, default = 0)
201 | 
202 |         parser.add_argument('--noise_iters', type = int, default = 500)
203 |         parser.add_argument('--noise_eps', type = float, default = 1e-3)
204 |         parser.add_argument('--noise_AMSGrad', type = bool, default = True)
205 |         parser.add_argument('--noise_learning_rate', type = float, default = 1e-1)
206 |         parser.add_argument('--noise_lambda_sourcelayer', type = float, default = 1e-3)
207 |         parser.add_argument('--noise_decrease_LR', type = int, default = 20)
208 |         parser.add_argument('--noise_sourceLayer', type = str, default = 'ReLU2')
209 |         parser.add_argument('--noise_targetLayer', type = str, default = 'fc3')
210 |         parser.add_argument('--noise_level', type = float, default = None)
211 |         parser.add_argument('--noise_epochs', type = int, default = 1)
212 |         parser.add_argument('--noise_batch_size', type = int, default = 32)
213 | 
214 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
215 |         parser.set_defaults(gpu=True)
216 | 
217 |         parser.add_argument('--novalidation', dest='validation', action='store_false')
218 |         parser.set_defaults(validation=True)
219 |         args = parser.parse_args()
220 | 
221 |         args.model_dir = "checkpoints/" + args.dataset + '/'
222 |         args.model_name = "ckpt.pth"
223 | 
224 |         if args.noise_level == None:
225 |             for nl in np.concatenate((np.arange(0, 110, 10), np.arange(100, 1100, 100)), axis=0):
226 |                 args.noise_level = nl
227 |                 noise_gen(
228 |                     args = args,
229 |                     model_dir = args.model_dir,
230 |                     model_name = args.model_name
231 |                 )
232 |         else:
233 |             noise_gen(
234 |                 args = args,
235 |                 model_dir = args.model_dir,
236 |                 model_name = args.model_name
237 |             )
238 | 
239 |     except:
240 |         traceback.print_exc(file=sys.stdout)
241 |         sys.exit(1)
242 | 


--------------------------------------------------------------------------------
/training.py:
--------------------------------------------------------------------------------
  1 | # @Author: zechenghe
  2 | # @Date:   2019-01-20T16:46:24-05:00
  3 | # @Last modified by:   zechenghe
  4 | # @Last modified time: 2019-02-01T14:01:19-05:00
  5 | 
  6 | import time
  7 | import math
  8 | import os
  9 | import numpy as np
 10 | 
 11 | import torch
 12 | import torchvision
 13 | import torchvision.transforms as transforms
 14 | import torch.nn as nn
 15 | import torch.nn.functional as F
 16 | import torch.optim as optim
 17 | import torch.backends.cudnn as cudnn
 18 | 
 19 | from net import *
 20 | from utils import *
 21 | 
 22 | def train(DATASET = 'CIFAR10', network = 'CIFAR10CNN', NEpochs = 200, imageWidth = 32,
 23 |         imageHeight = 32, imageSize = 32*32, NChannels = 3, NClasses = 10,
 24 |         BatchSize = 32, learningRate = 1e-3, NDecreaseLR = 20, eps = 1e-3,
 25 |         AMSGrad = True, model_dir = "checkpoints/CIFAR10/", model_name = "ckpt.pth", gpu = True):
 26 | 
 27 |     print "DATASET: ", DATASET
 28 | 
 29 |     if DATASET == 'MNIST':
 30 | 
 31 |         mu = torch.tensor([0.5], dtype=torch.float32)
 32 |         sigma = torch.tensor([0.5], dtype=torch.float32)
 33 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 34 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 35 | 
 36 |         tsf = {
 37 |             'train': transforms.Compose(
 38 |             [
 39 |                 transforms.ToTensor(),
 40 |                 Normalize
 41 |             ]),
 42 | 
 43 |             'test': transforms.Compose(
 44 |             [
 45 |                 transforms.ToTensor(),
 46 |                 Normalize
 47 |             ])
 48 |         }
 49 | 
 50 |         trainset = torchvision.datasets.MNIST(root='./data/MNIST', train=True,
 51 |                                         download=True, transform = tsf['train'])
 52 | 
 53 |         testset = torchvision.datasets.MNIST(root='./data/MNIST', train=False,
 54 |                                        download=True, transform = tsf['test'])
 55 | 
 56 |     elif DATASET == 'CIFAR10':
 57 | 
 58 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
 59 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
 60 |         Normalize = transforms.Normalize(mu.tolist(), sigma.tolist())
 61 |         Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 62 |         tsf = {
 63 |             'train': transforms.Compose(
 64 |             [
 65 |             transforms.RandomHorizontalFlip(),
 66 |             transforms.RandomAffine(degrees = 10, translate = [0.1, 0.1], scale = [0.9, 1.1]),
 67 |             transforms.ToTensor(),
 68 |             Normalize
 69 |             ]),
 70 |             'test': transforms.Compose(
 71 |             [
 72 |             transforms.ToTensor(),
 73 |             Normalize
 74 |             ])
 75 |         }
 76 |         trainset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = True,
 77 |                                         download=True, transform = tsf['train'])
 78 |         testset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train = False,
 79 |                                        download=True, transform = tsf['test'])
 80 | 
 81 |     netDict = {
 82 |         'LeNet': LeNet,
 83 |         'CIFAR10CNN': CIFAR10CNN
 84 |     }
 85 | 
 86 |     if network in netDict:
 87 |         net = netDict[network](NChannels)
 88 |     else:
 89 |         print "Network not found"
 90 |         exit(1)
 91 | 
 92 |     print net
 93 |     print "len(trainset) ", len(trainset)
 94 |     print "len(testset) ", len(testset)
 95 |     x_train, y_train = trainset.data, trainset.targets,
 96 |     x_test, y_test = testset.data, testset.targets,
 97 | 
 98 |     print "x_train.shape ", x_train.shape
 99 |     print "x_test.shape ", x_test.shape
100 | 
101 |     trainloader = torch.utils.data.DataLoader(trainset, batch_size = BatchSize,
102 |                                       shuffle = True, num_workers = 1)
103 |     testloader = torch.utils.data.DataLoader(testset, batch_size = 1000,
104 |                                       shuffle = False, num_workers = 1)
105 |     trainIter = iter(trainloader)
106 |     testIter = iter(testloader)
107 | 
108 |     criterion = nn.CrossEntropyLoss()
109 |     softmax = nn.Softmax(dim=1)
110 | 
111 |     if gpu:
112 |         net.cuda()
113 |         criterion.cuda()
114 |         softmax.cuda()
115 | 
116 |     optimizer = optim.Adam(params = net.parameters(), lr = learningRate, eps = eps, amsgrad = AMSGrad)
117 | 
118 |     NBatch = len(trainset) / BatchSize
119 |     cudnn.benchmark = True
120 |     for epoch in range(NEpochs):
121 |         lossTrain = 0.0
122 |         accTrain = 0.0
123 |         for i in range(NBatch):
124 |             try:
125 |                 batchX, batchY = trainIter.next()
126 |             except StopIteration:
127 |                 trainIter = iter(trainloader)
128 |                 batchX, batchY = trainIter.next()
129 | 
130 |             if gpu:
131 |                 batchX = batchX.cuda()
132 |                 batchY = batchY.cuda()
133 | 
134 |             optimizer.zero_grad()
135 |             logits = net.forward(batchX)
136 |             prob = softmax(logits)
137 | 
138 |             loss = criterion(logits, batchY)
139 |             loss.backward()
140 |             optimizer.step()
141 | 
142 |             lossTrain += loss.cpu().detach().numpy() / NBatch
143 |             if gpu:
144 |                 pred = np.argmax(prob.cpu().detach().numpy(), axis = 1)
145 |                 groundTruth = batchY.cpu().detach().numpy()
146 |             else:
147 |                 pred = np.argmax(prob.detach().numpy(), axis = 1)
148 |                 groundTruth = batchY.detach().numpy()
149 | 
150 |             acc = np.mean(pred == groundTruth)
151 |             accTrain += acc / NBatch
152 | 
153 |         if (epoch + 1) % NDecreaseLR == 0:
154 |             learningRate = learningRate / 2.0
155 |             setLearningRate(optimizer, learningRate)
156 | 
157 |         print "Epoch: ", epoch, "Loss: ", lossTrain, "Train accuracy: ", accTrain
158 | 
159 |         accTest = evalTest(testloader, net, gpu = gpu)
160 | 
161 |     if not os.path.exists(model_dir):
162 |         os.makedirs(model_dir)
163 |     torch.save(net, model_dir + model_name)
164 |     print "Model saved"
165 | 
166 |     newNet = torch.load(model_dir + model_name)
167 |     newNet.eval()
168 |     accTest = evalTest(testloader, net, gpu = gpu)
169 |     print "Model restore done"
170 | 
171 | 
172 | if __name__ == '__main__':
173 |     import argparse
174 |     import sys
175 |     import traceback
176 | 
177 |     try:
178 |         parser = argparse.ArgumentParser()
179 |         parser.add_argument('--dataset', type = str, default = 'CIFAR10')
180 |         parser.add_argument('--network', type = str, default = 'CIFAR10CNN')
181 |         parser.add_argument('--epochs', type = int, default = 200)
182 |         parser.add_argument('--eps', type = float, default = 1e-3)
183 |         parser.add_argument('--AMSGrad', type = bool, default = True)
184 |         parser.add_argument('--batch_size', type = int, default = 32)
185 |         parser.add_argument('--learning_rate', type = float, default = 1e-3)
186 |         parser.add_argument('--decrease_LR', type = int, default = 20)
187 | 
188 |         parser.add_argument('--nogpu', dest='gpu', action='store_false')
189 |         parser.set_defaults(gpu=True)
190 |         args = parser.parse_args()
191 | 
192 |         model_dir = "checkpoints/" + args.dataset + '/'
193 |         model_name = "ckpt.pth"
194 | 
195 |         if args.dataset == 'MNIST':
196 | 
197 |             imageWidth = 28
198 |             imageHeight = 28
199 |             imageSize = imageWidth * imageHeight
200 |             NChannels = 1
201 |             NClasses = 10
202 |             network = 'LeNet'
203 | 
204 |         elif args.dataset == 'CIFAR10':
205 | 
206 |             imageWidth = 32
207 |             imageHeight = 32
208 |             imageSize = imageWidth * imageHeight
209 |             NChannels = 3
210 |             NClasses = 10
211 |             network = 'CIFAR10CNN'
212 | 
213 |         else:
214 |             print "No Dataset Found"
215 |             exit(0)
216 | 
217 |         train(DATASET = args.dataset, network = network, NEpochs = args.epochs, imageWidth = imageWidth,
218 |         imageHeight = imageHeight, imageSize = imageSize, NChannels = NChannels, NClasses = NClasses,
219 |         BatchSize = args.batch_size, learningRate = args.learning_rate, NDecreaseLR = args.decrease_LR, eps = args.eps,
220 |         AMSGrad = args.AMSGrad, model_dir = model_dir, model_name = model_name, gpu = args.gpu)
221 | 
222 |     except:
223 |         traceback.print_exc(file=sys.stdout)
224 |         sys.exit(1)
225 | 


--------------------------------------------------------------------------------
/utils.py:
--------------------------------------------------------------------------------
  1 | # @Author: zechenghe
  2 | # @Date:   2019-01-21T12:01:09-05:00
  3 | # @Last modified by:   zechenghe
  4 | # @Last modified time: 2019-02-01T14:50:41-05:00
  5 | 
  6 | import time
  7 | import math
  8 | import os
  9 | import numpy as np
 10 | 
 11 | import torch
 12 | import torchvision
 13 | import torchvision.transforms as transforms
 14 | import torch.nn as nn
 15 | import torch.nn.functional as F
 16 | import torch.optim as optim
 17 | import torch.backends.cudnn as cudnn
 18 | 
 19 | 
 20 | def accuracy(predictions, labels):
 21 | 
 22 |     if not (predictions.shape == labels.shape):
 23 |         print "predictions.shape ", predictions.shape, "labels.shape ", labels.shape
 24 |         raise AssertionError
 25 | 
 26 |     correctly_predicted = np.sum(np.argmax(predictions, 1) == np.argmax(labels, 1))
 27 |     accu = (100.0 * correctly_predicted) / predictions.shape[0]
 28 |     return accu
 29 | 
 30 | def pseudoInverse(W):
 31 |     return np.linalg.pinv(W)
 32 | 
 33 | 
 34 | def getImgByClass(Itr, C = None):
 35 | 
 36 |     if C == None:
 37 |         return Itr.next()
 38 | 
 39 |     while (True):
 40 |         img, label = Itr.next()
 41 |         if label == C:
 42 |             break
 43 |     return img, label
 44 | 
 45 | 
 46 | def clip(data):
 47 |     data[data > 1.0] = 1.0
 48 |     data[data < 0.0] = 0.0
 49 |     return data
 50 | 
 51 | def preprocess(data):
 52 | 
 53 |     size = data.shape
 54 |     NChannels = size[-1]
 55 |     assert NChannels == 1 or NChannels == 3
 56 |     if NChannels == 1:
 57 |         mu = 0.5
 58 |         sigma = 0.5
 59 |     elif NChannels == 3:
 60 |         mu = [0.485, 0.456, 0.406]
 61 |         sigma = [0.229, 0.224, 0.225]
 62 |     data = (data - mu) / sigma
 63 | 
 64 |     assert data.shape == size
 65 |     return data
 66 | 
 67 | 
 68 | def deprocess(data):
 69 | 
 70 |     assert len(data.size()) == 4
 71 | 
 72 |     BatchSize = data.size()[0]
 73 |     assert BatchSize == 1
 74 | 
 75 |     NChannels = data.size()[1]
 76 |     if NChannels == 1:
 77 |         mu = torch.tensor([0.5], dtype=torch.float32)
 78 |         sigma = torch.tensor([0.5], dtype=torch.float32)
 79 |     elif NChannels == 3:
 80 |         mu = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)
 81 |         sigma = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)
 82 |     else:
 83 |         print "Unsupported image in deprocess()"
 84 |         exit(1)
 85 | 
 86 |     Unnormalize = transforms.Normalize((-mu / sigma).tolist(), (1.0 / sigma).tolist())
 87 |     return clip(Unnormalize(data[0,:,:,:]).unsqueeze(0))
 88 | 
 89 | 
 90 | def evalTest(testloader, net, gpu = True):
 91 |     testIter = iter(testloader)
 92 |     acc = 0.0
 93 |     NBatch = 0
 94 |     for i, data in enumerate(testIter, 0):
 95 |         NBatch += 1
 96 |         batchX, batchY = data
 97 |         if gpu:
 98 |             batchX = batchX.cuda()
 99 |             batchY = batchY.cuda()
100 |         logits = net.forward(batchX)
101 | 
102 |         if gpu:
103 |             pred = np.argmax(logits.cpu().detach().numpy(), axis = 1)
104 |             groundTruth = batchY.cpu().detach().numpy()
105 |         else:
106 |             pred = np.argmax(logits.detach().numpy(), axis = 1)
107 |             groundTruth = batchY.detach().numpy()
108 |         acc += np.mean(pred == groundTruth)
109 |     accTest = acc / NBatch
110 |     print "Test accuracy: ", accTest #, "NBatch: ", NBatch, "pred == groundTruth.shape", (pred == groundTruth).shape
111 |     return accTest
112 | 
113 | 
114 | 
115 | def weight_init(m):
116 |     '''
117 |     Usage:
118 |         model = Model()
119 |         model.apply(weight_init)
120 |     '''
121 |     if isinstance(m, nn.Conv1d):
122 |         init.normal_(m.weight.data)
123 |         if m.bias is not None:
124 |             init.normal_(m.bias.data)
125 |     elif isinstance(m, nn.Conv2d):
126 |         init.xavier_normal_(m.weight.data)
127 |         if m.bias is not None:
128 |             init.normal_(m.bias.data)
129 |     elif isinstance(m, nn.Conv3d):
130 |         init.xavier_normal_(m.weight.data)
131 |         if m.bias is not None:
132 |             init.normal_(m.bias.data)
133 |     elif isinstance(m, nn.ConvTranspose1d):
134 |         init.normal_(m.weight.data)
135 |         if m.bias is not None:
136 |             init.normal_(m.bias.data)
137 |     elif isinstance(m, nn.ConvTranspose2d):
138 |         init.xavier_normal_(m.weight.data)
139 |         if m.bias is not None:
140 |             init.normal_(m.bias.data)
141 |     elif isinstance(m, nn.ConvTranspose3d):
142 |         init.xavier_normal_(m.weight.data)
143 |         if m.bias is not None:
144 |             init.normal_(m.bias.data)
145 |     elif isinstance(m, nn.BatchNorm1d):
146 |         init.normal_(m.weight.data, mean=1, std=0.02)
147 |         init.constant_(m.bias.data, 0)
148 |     elif isinstance(m, nn.BatchNorm2d):
149 |         init.normal_(m.weight.data, mean=1, std=0.02)
150 |         init.constant_(m.bias.data, 0)
151 |     elif isinstance(m, nn.BatchNorm3d):
152 |         init.normal_(m.weight.data, mean=1, std=0.02)
153 |         init.constant_(m.bias.data, 0)
154 |     elif isinstance(m, nn.Linear):
155 |         init.xavier_normal_(m.weight.data)
156 |         init.normal_(m.bias.data)
157 |     elif isinstance(m, nn.LSTM):
158 |         for param in m.parameters():
159 |             if len(param.shape) >= 2:
160 |                 init.orthogonal_(param.data)
161 |             else:
162 |                 init.normal_(param.data)
163 |     elif isinstance(m, nn.LSTMCell):
164 |         for param in m.parameters():
165 |             if len(param.shape) >= 2:
166 |                 init.orthogonal_(param.data)
167 |             else:
168 |                 init.normal_(param.data)
169 |     elif isinstance(m, nn.GRU):
170 |         for param in m.parameters():
171 |             if len(param.shape) >= 2:
172 |                 init.orthogonal_(param.data)
173 |             else:
174 |                 init.normal_(param.data)
175 |     elif isinstance(m, nn.GRUCell):
176 |         for param in m.parameters():
177 |             if len(param.shape) >= 2:
178 |                 init.orthogonal_(param.data)
179 |             else:
180 |                 init.normal_(param.data)
181 | 
182 | def setLearningRate(optimizer, lr):
183 |     """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
184 |     for param_group in optimizer.param_groups:
185 |         param_group['lr'] = lr
186 | 
187 | 
188 | def TV(x):
189 |     batch_size = x.size()[0]
190 |     h_x = x.size()[2]
191 |     w_x = x.size()[3]
192 |     count_h = _tensor_size(x[:,:,1:,:])
193 |     count_w = _tensor_size(x[:,:,:,1:])
194 |     h_tv = torch.pow(x[:,:,1:,:]-x[:,:,:h_x-1,:], 2).sum()
195 |     w_tv = torch.pow(x[:,:,:,1:]-x[:,:,:,:w_x-1], 2).sum()
196 |     return (h_tv / count_h + w_tv / count_w) / batch_size
197 | 
198 | def _tensor_size(t):
199 |     return t.size()[1]*t.size()[2]*t.size()[3]
200 | 
201 | 
202 | def l2loss(x):
203 |     return (x**2).mean()
204 | 
205 | def l1loss(x):
206 |     return (torch.abs(x)).mean()
207 | 
208 | def getL1Stat(net, x):
209 |     for layer in net.layerDict:
210 |         targetLayer = net.layerDict[layer]
211 |         layerOutput = net.getLayerOutput(x, targetLayer)
212 |         print "Layer " + layer + ' l1 loss:', l1loss(layerOutput).cpu().detach().numpy()
213 | 
214 | 
215 | def getModule(net, blob):
216 |     modules = blob.split('.')
217 | #    print "Target layer: ", modules
218 | #    if len(modules) == 1:
219 | #        return net._modules.get(blob)
220 | #    else:
221 | 
222 |     curr_module = net
223 |     print curr_module
224 |     for m in modules:
225 |         curr_module = curr_module._modules.get(m)
226 |     return curr_module
227 | 
228 | def getLayerOutputHook(module, input, output):
229 |     if not hasattr(module, 'activations'):
230 |         module.activations = []
231 |     module.activations.append(output)
232 | 
233 | def getHookActs(model, module, input):
234 |     if hasattr(module, 'activations'):
235 |         del module.activations[:]
236 |     _ = model.forward(input)
237 |     assert(len(module.activations) == 1)
238 |     return module.activations[0]
239 | 
240 | def saveImage(img, filepath):
241 |     torchvision.utils.save_image(img, filepath)
242 | 
243 | 
244 | def apply_noise(input, noise_type, noise_level, mean=0.0, gpu=True, args=None):
245 | 
246 |     if noise_type == 'Gaussian':
247 |         noise = torch.randn(input.size()) * noise_level + mean
248 |         noise = noise.cuda() if gpu else noise
249 |         output = input + noise
250 | 
251 |     elif noise_type == 'Laplace':
252 |         noise = np.random.laplace(
253 |             loc= mean,
254 |             scale = noise_level,
255 |             size = input.size()
256 |         )
257 |         noise = torch.tensor(noise, dtype = torch.float)
258 |         noise = noise.cuda() if gpu else noise
259 |         output = input + noise
260 | 
261 |     elif noise_type == 'dropout':
262 |         mask = np.random.choice([0.0, 1.0], size=input.size(), replace=True, p=[noise_level, 1-noise_level])
263 |         mask = torch.tensor(mask, dtype = torch.float)
264 |         mask = mask.cuda() if gpu else mask
265 |         output = input * mask
266 | 
267 |     elif noise_type == 'dropout-non-zero':
268 |         input_list = input.detach().cpu().numpy().reshape([-1])
269 |         output = input_list.copy()
270 | 
271 |         for i in range(len(input_list)):
272 |             if input_list[i] > 0:
273 |                 if np.random.rand() < noise_level:
274 |                     output[i] = -1.0
275 |             else:
276 |                 output[i] = -np.random.rand() * 10.0
277 |         output = torch.tensor(np.array(output).reshape(input.size()), dtype = torch.float)
278 |         output = output.cuda() if gpu else output
279 | 
280 |     elif noise_type == 'redistribute':
281 |         input_list = input.detach().cpu().numpy().reshape([-1])
282 |         idx = np.argsort(input_list)
283 |         map = np.linspace(start=0.0, stop=1.0, num=len(input_list))
284 | 
285 |         output = [0]*len(input_list)
286 |         for i in range(len(idx)):
287 |             if input_list[idx[i]] != 0 and np.random.rand() > noise_level:
288 |                 output[idx[i]] = 1.0
289 |         output = torch.tensor(np.array(output).reshape(input.size()), dtype = torch.float)
290 |         output = output.cuda() if gpu else output
291 | 
292 |         #print "input", input
293 |         #print "output", output
294 | 
295 |     elif noise_type == 'impulse':
296 |         noise = np.random.choice([0.0, 1.0], size=input.size(), replace=True, p=[1-noise_level, noise_level])
297 |         noise = torch.tensor(noise, dtype = torch.float)
298 |         noise = noise.cuda() if gpu else noise
299 |         output = input + noise
300 | 
301 |     elif noise_type == 'noise_gen' or 'noise_gen_opt':
302 |         noise_dir = 'noise' + ('_opt' if noise_type == 'noise_gen_opt' else "") + '/' + args.dataset + '/'
303 |         noise_file_name = args.noise_sourceLayer + '-' + args.noise_targetLayer + '-' + str(round(noise_level, 2))
304 | 
305 |         noise = np.load(noise_dir + noise_file_name + '.npy')
306 |         noise = torch.tensor(noise, dtype = torch.float)
307 | 
308 |         batch_size = input.size()[0]
309 |         noise = torch.cat(batch_size * [noise])
310 |         noise = noise.cuda() if gpu else noise
311 |         output = input + noise
312 | 
313 |     else:
314 |         print "Unsupported Noise Type: ", noise_type
315 |         exit(1)
316 | 
317 |     return output
318 | 
319 | def evalTestSplitModel(testloader, netEdge, netCloud, layer, gpu, noise_type = None, noise_level = 0.0, args=None):
320 |     testIter = iter(testloader)
321 |     acc = 0.0
322 |     NBatch = 0
323 |     for i, data in enumerate(testIter, 0):
324 |         batchX, batchY = data
325 |         if gpu:
326 |             batchX = batchX.cuda()
327 |             batchY = batchY.cuda()
328 | 
329 |         if hasattr(args, 'add_noise_to_input') and args.add_noise_to_input:
330 |             batchX = apply_noise(batchX, noise_type, noise_level, gpu=gpu, args=args)
331 | 
332 |         try:
333 |             edgeOutput = netEdge.getLayerOutput(batchX, netEdge.layerDict[layer]).clone()
334 |         except Exception, e:
335 |             #print "Except in evalTestSplitModel getLayerOutput, this is a Edge-only model"
336 |             #print str(e)
337 |             edgeOutput = netEdge.forward(batchX).clone()
338 | 
339 |         if noise_type != None and not (hasattr(args, 'add_noise_to_input') and args.add_noise_to_input):
340 |             edgeOutput = apply_noise(edgeOutput, noise_type, noise_level, gpu=gpu, args=args)
341 | 
342 |         #cloudOuput = net.forward(batchX)
343 |         logits = netCloud.forward_from(edgeOutput, layer)
344 | 
345 |         #softmax = nn.Softmax().cuda()
346 |         #prob = softmax(logits)
347 |         #print prob[:100,:].max(dim=1)
348 | 
349 |         if gpu:
350 |             pred = np.argmax(logits.cpu().detach().numpy(), axis = 1)
351 |             groundTruth = batchY.cpu().detach().numpy()
352 |         else:
353 |             pred = np.argmax(logits.detach().numpy(), axis = 1)
354 |             groundTruth = batchY.detach().numpy()
355 |         acc += np.mean(pred == groundTruth)
356 |         NBatch += 1
357 | 
358 |     accTest = acc / NBatch
359 |     #print "Test accuracy: ", accTest #, "NBatch: ", NBatch, "pred == groundTruth.shape", (pred == groundTruth).shape
360 |     return accTest
361 | 
362 | def get_PSNR(refimg, invimg, peak = 1.0):
363 |     psnr = 10*np.log10(peak**2 / np.mean((refimg - invimg)**2))
364 |     return psnr
365 | 


--------------------------------------------------------------------------------
/utils.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/zechenghe/Inverse_Collaborative_Inference/f010aed0035238c573037b04c47c334ea733d15c/utils.pyc


--------------------------------------------------------------------------------