├── README.md ├── abs_test.py ├── data ├── npy_data.py └── utils.py ├── generate_examples.py ├── get_data.py ├── models ├── nin.py ├── resnet.py ├── resnet_more_activation.py ├── vgg.py └── vgg_trojai.py ├── none ├── compromised_detection.py ├── reset_neurons.py └── test_utils.py ├── requirements.txt ├── run_none.py └── test_natral_trojan.py /README.md: -------------------------------------------------------------------------------- 1 | # NONE 2 | Code for "Training with More Confidence: Mitigating Injected and Natural Backdoors During Training" 3 | 4 | ## Environment 5 | requirements.txt 6 | 7 | ## Prepare Data 8 | - We provide poisonsed datasets: 9 | 10 | https://drive.google.com/drive/folders/13oBVDijC-nweBtHw2CBbcevXemc5vTTW. 11 | 12 | - Download and unzip them: 13 | ```bash 14 | unzip single_target_p5_s3.zip 15 | ``` 16 | - You can also generate poisoned datasets by using https://github.com/MadryLab/label-consistent-backdoor-code 17 | 18 | ## Run Experiments on Defending Injected Trojans 19 | - For example, to run NONE on CIFAR10 against BadNets Attack: 20 | 21 | ```bash 22 | python -u run_none.py --dataset cifar10 --arch resnet18 --poison_type single_target --none_lr 1e-4 --max_reset_fraction 0.03 --poison_rate 0.05 --epoch_num_1 200 --epoch_num_2 20 23 | ``` 24 | 25 | ## Run Experiments on Defending Natural Trojans 26 | 27 | - First generate examples for calculating REASR: 28 | ```bash 29 | python generate_examples.py --dataset cifar10 --save_dir ./abs/example/ 30 | ``` 31 | 32 | - Run NONE to against Natural Trojans on CIFAR10: 33 | 34 | ```bash 35 | python -u run_none.py --dataset cifar10 --arch nin --poison_type none --none_lr 1e-4 --max_reset_fraction 0.20 --epoch_num_1 200 --epoch_num_2 40 --round_num 10 36 | ``` 37 | 38 | - The results will be recorded in logger file. 39 | -------------------------------------------------------------------------------- /abs_test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import argparse 4 | 5 | import sys 6 | import json 7 | import skimage.io 8 | import random 9 | import torch 10 | import torch.nn.functional as F 11 | import pickle 12 | import time 13 | import torch.nn as nn 14 | 15 | from models import nin,resnet,vgg,vgg_trojai 16 | 17 | np.set_printoptions(precision=2, linewidth=200, threshold=10000) 18 | 19 | def seed_torch(seed=333):#333 20 | 21 | random.seed(seed) 22 | os.environ['PYTHONHASHSEED'] = str(seed) 23 | np.random.seed(seed) 24 | torch.manual_seed(seed) 25 | torch.cuda.manual_seed(seed) 26 | torch.cuda.manual_seed_all(seed) 27 | torch.backends.cudnn.deterministic = True 28 | 29 | seed_torch() 30 | print(torch.rand(1)) 31 | print(torch.rand(1)) 32 | 33 | class Net(nn.Module): 34 | def __init__(self): 35 | super(Net, self).__init__() 36 | self.classifier = nn.Sequential( 37 | nn.Conv2d(3, 192, kernel_size=5, stride=1, padding=2), 38 | nn.ReLU(inplace=True), 39 | nn.Conv2d(192, 160, kernel_size=1, stride=1, padding=0), 40 | nn.ReLU(inplace=True), 41 | nn.Conv2d(160, 96, kernel_size=1, stride=1, padding=0), 42 | nn.ReLU(inplace=True), 43 | nn.MaxPool2d(kernel_size=3, stride=2, padding=1), 44 | nn.Dropout(0.5), 45 | 46 | nn.Conv2d(96, 192, kernel_size=5, stride=1, padding=2), 47 | nn.ReLU(inplace=True), 48 | nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0), 49 | nn.ReLU(inplace=True), 50 | nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0), 51 | nn.ReLU(inplace=True), 52 | nn.AvgPool2d(kernel_size=3, stride=2, padding=1), 53 | nn.Dropout(0.5), 54 | 55 | nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1), 56 | nn.ReLU(inplace=True), 57 | nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0), 58 | nn.ReLU(inplace=True), 59 | nn.Conv2d(192, 10, kernel_size=1, stride=1, padding=0), 60 | nn.ReLU(inplace=True), 61 | nn.AvgPool2d(kernel_size=8, stride=1, padding=0), 62 | 63 | ) 64 | 65 | def forward(self, x): 66 | x = self.classifier(x) 67 | x = x.view(x.size(0), 10) 68 | return x 69 | 70 | 71 | 72 | # with open('./config.json') as config_file: 73 | # config = json.load(config_file) 74 | config = {} 75 | config['gpu_id'] = '0' 76 | config['print_level'] = 2 77 | config['random_seed'] = 333 78 | config['num_classes'] = 10 79 | config['channel_last'] = 0 80 | config['w'] = 32 81 | config['h'] = 32 82 | config['reasr_bound'] = 0.2 83 | config['batch_size'] = 20 84 | config['has_softmax'] = 0 85 | config['samp_k'] = 1 86 | #config['samp_k'] = 5 87 | config['same_range'] = 0 88 | config['n_samples'] = 3 89 | config['samp_batch_size'] = 8 90 | config['top_n_neurons'] = 10 91 | #config['top_n_neurons'] = 100 92 | config['re_batch_size'] = 80 93 | #config['max_troj_size'] = 64 94 | config['max_troj_size'] = 64 95 | config['filter_multi_start'] = 1 96 | config['re_mask_lr'] = 0.1 97 | #config['re_mask_lr'] = 5e-1 98 | config['re_mask_weight'] = 50 99 | config['mask_multi_start'] = 1 100 | config['re_epochs'] = 50 101 | config['n_re_samples'] = 240 102 | 103 | #score_type = "max" 104 | score_type = "mean_exclude_target" 105 | 106 | #config['re_epochs'] = 50 107 | #config['n_re_samples'] = 10000 108 | 109 | #config['n_re_samples'] = 1000 110 | 111 | #config['re_epochs'] = 100 112 | #config['n_re_samples'] = 10000 113 | 114 | 115 | channel_last = bool(config['channel_last']) 116 | random_seed = int(config['random_seed']) 117 | os.environ["CUDA_VISIBLE_DEVICES"] = config["gpu_id"] 118 | 119 | #np.random.seed(random_seed) 120 | #seed_torch(random_seed) 121 | w = config["w"] 122 | h = config["h"] 123 | num_classes = config["num_classes"] 124 | use_mask = True 125 | count_mask = True 126 | tdname = 'temp' 127 | window_size = 12 128 | mask_epsilon = 0.01 129 | delta_shape = [window_size,window_size,3,3] 130 | Troj_size = config['max_troj_size'] 131 | reasr_bound = float(config['reasr_bound']) 132 | top_n_neurons = int(config['top_n_neurons']) 133 | mask_multi_start = int(config['mask_multi_start']) 134 | filter_multi_start = int(config['filter_multi_start']) 135 | re_mask_weight = float(config['re_mask_weight']) 136 | re_mask_lr = float(config['re_mask_lr']) 137 | batch_size = config['batch_size'] 138 | has_softmax = bool(config['has_softmax']) 139 | print('channel_last', channel_last, 'gpu_id', config["gpu_id"], 'has softmax', has_softmax) 140 | # try_epochs = 10 141 | 142 | # cifar= CIFAR10() 143 | # print('gpu id', config["gpu_id"]) 144 | # l_bounds = cifar.l_bounds 145 | # h_bounds = cifar.h_bounds 146 | # print('mean', cifar.mean, 'std', cifar.std, 'l bounds', l_bounds[0,0,0], 'h_bounds', h_bounds[0,0,0]) 147 | 148 | # l_bounds_channel_first = np.transpose(l_bounds, [0,3,1,2]) 149 | # h_bounds_channel_first = np.transpose(h_bounds, [0,3,1,2]) 150 | Print_Level = int(config['print_level']) 151 | re_epochs = int(config['re_epochs']) 152 | n_re_samples = int(config['n_re_samples']) 153 | 154 | arch = None 155 | 156 | def preprocess(img): 157 | img = np.transpose(img, [0,3,1,2]) 158 | return img.astype(np.float32) / 255.0 159 | 160 | '''def deprocess(img): 161 | img = np.transpose(img, [0,2,3,1]) 162 | return (img*255).astype(np.uint8)''' 163 | 164 | def deprocess(img): 165 | mean = [0.4914, 0.4822, 0.4465] 166 | std = [0.2023, 0.1994, 0.2010] 167 | img_copy = img.copy() 168 | for j in range(img_copy.shape[0]): 169 | for i in range(len(mean)): 170 | #img[i] = (img[i] - mean[i])/std[i] 171 | img_copy[j][i] = img_copy[j][i]*std[i] + mean[i] 172 | img_copy = np.transpose(img_copy, [0,2,3,1]) 173 | return (img_copy*255).astype(np.uint8) 174 | 175 | 176 | def check_values(images, labels, model, children, target_layers): 177 | maxes = {} 178 | print("check_values") 179 | #print(images) 180 | for layer_i in range(0, len(children) - 1): 181 | if not children[layer_i].__class__.__name__ in target_layers: 182 | continue 183 | temp_model1 = torch.nn.Sequential(*children[:layer_i+1]) 184 | 185 | max_val = -np.inf 186 | print(len(images)) 187 | print(batch_size) 188 | for i in range(len(images)//batch_size): 189 | batch_data = torch.FloatTensor(images[batch_size*i:batch_size*(i+1)]) 190 | batch_data = batch_data.cuda() 191 | inner_outputs = temp_model1(batch_data).cpu().detach().numpy() 192 | if channel_last: 193 | n_neurons = inner_outputs.shape[-1] 194 | else: 195 | n_neurons = inner_outputs.shape[1] 196 | max_val = np.maximum(max_val, np.amax(inner_outputs)) 197 | print(np.amax(inner_outputs)) 198 | 199 | key = '{0}_{1}'.format(children[layer_i].__class__.__name__, layer_i) 200 | maxes[key] = [max_val] 201 | print('max val', key, max_val) 202 | del temp_model1, batch_data, inner_outputs 203 | return maxes 204 | 205 | def sample_neuron(images, labels, model, children, target_layers, model_type, mvs): 206 | all_ps = {} 207 | samp_k = config['samp_k'] 208 | same_range = config['same_range'] 209 | n_samples = config['n_samples'] 210 | sample_batch_size = config['samp_batch_size'] 211 | if model_type == 'DenseNet': 212 | sample_batch_size = max(sample_batch_size // 3, 1) 213 | n_images = images.shape[0] 214 | if Print_Level > 0: 215 | print('sampling n imgs', n_images) 216 | 217 | end_layer = len(children)-1 218 | if has_softmax: 219 | end_layer = len(children)-2 220 | 221 | # maxv_layer_i = 0 222 | # fmaxv = -np.inf 223 | # for layer_i in range(2, end_layer): 224 | # if not children[layer_i].__class__.__name__ in target_layers: 225 | # continue 226 | # mv_key = '{0}_{1}'.format(children[layer_i].__class__.__name__, layer_i) 227 | # maxv = max(mvs[mv_key]) 228 | # if maxv > fmaxv: 229 | # fmaxv = maxv 230 | # maxv_layer_i = layer_i 231 | # sample_layers = [maxv_layer_i] 232 | 233 | sample_layers = [] 234 | for layer_i in range(2, end_layer): 235 | if not children[layer_i].__class__.__name__ in target_layers: 236 | continue 237 | sample_layers.append(layer_i) 238 | #sample_layers = sample_layers[-3:-1] 239 | if arch == "nin": 240 | sample_layers = sample_layers[-2:-1] 241 | elif arch == "lenet": 242 | sample_layers = sample_layers[-2:-1] 243 | elif arch == "vgg16": 244 | sample_layers = sample_layers[-1:] 245 | elif arch == "vgg11": 246 | sample_layers = sample_layers[-1:] 247 | else: 248 | sample_layers = sample_layers[-1:] 249 | 250 | #sample_layers = sample_layers[-2:-1] 251 | 252 | # sample_layers = sample_layers 253 | 254 | for layer_i in sample_layers: 255 | #print("layer_i",layer_i) 256 | #if Print_Level > 0: 257 | #print('layer', layer_i, children[layer_i]) 258 | temp_model1 = torch.nn.Sequential(*children[:layer_i+1]) 259 | if has_softmax: 260 | temp_model2 = torch.nn.Sequential(*children[layer_i+1:-1]) 261 | else: 262 | temp_model2 = torch.nn.Sequential(*children[layer_i+1:]) 263 | 264 | if same_range: 265 | vs = np.asarray([i*samp_k for i in range(n_samples)]) 266 | else: 267 | mv_key = '{0}_{1}'.format(children[layer_i].__class__.__name__, layer_i) 268 | 269 | # tr = samp_k * max(mvs[mv_key])/(n_samples - 1) 270 | # vs = np.asarray([i*tr for i in range(n_samples)]) 271 | 272 | maxv = max(mvs[mv_key]) 273 | e_scale = np.array([0] + [np.power(2., i-1) for i in range(n_samples)]) 274 | vs = max(mvs[mv_key]) * e_scale 275 | 276 | #print('mv_key', vs) 277 | 278 | for input_i in range(n_images//batch_size): 279 | #print("input_i",input_i) 280 | batch_data = torch.FloatTensor(images[batch_size*input_i:batch_size*(input_i+1)]) 281 | batch_data = batch_data.cuda() 282 | inner_outputs = temp_model1(batch_data).cpu().detach().numpy() 283 | if channel_last: 284 | n_neurons = inner_outputs.shape[-1] 285 | else: 286 | n_neurons = inner_outputs.shape[1] 287 | 288 | nbatches = n_neurons//sample_batch_size 289 | #print(nbatches) 290 | for nt in range(nbatches): 291 | #print("nt",nt) 292 | l_h_t = [] 293 | for neuron in range(sample_batch_size): 294 | # h_t = torch.repeat_interleave(inner_outputs, repeats=n_samples, dim=0) 295 | if len(inner_outputs.shape) == 4: 296 | h_t = np.tile(inner_outputs, (n_samples, 1, 1, 1)) 297 | else: 298 | h_t = np.tile(inner_outputs, (n_samples, 1)) 299 | 300 | for i,v in enumerate(vs): 301 | if len(inner_outputs.shape) == 4: 302 | if channel_last: 303 | h_t[i*batch_size:(i+1)*batch_size,:,:,neuron+nt*sample_batch_size] = v 304 | else: 305 | h_t[i*batch_size:(i+1)*batch_size,neuron+nt*sample_batch_size,:,:] = v 306 | else: 307 | h_t[i*batch_size:(i+1)*batch_size, neuron+nt*sample_batch_size] = v 308 | l_h_t.append(h_t) 309 | # f_h_t = torch.cat(l_h_t, axis=0) 310 | f_h_t = np.concatenate(l_h_t, axis=0) 311 | 312 | f_h_t_t = torch.FloatTensor(f_h_t).cuda() 313 | print(temp_model2) 314 | print(f_h_t_t.shape) 315 | print(temp_model2[0](f_h_t_t).shape) 316 | print(temp_model2[:2](f_h_t_t).shape) 317 | print(temp_model2[:3](f_h_t_t).shape) 318 | fps = temp_model2(f_h_t_t).cpu().detach().numpy() 319 | # if Print_Level > 1: 320 | # print(nt, n_neurons, 'inner_outputs', inner_outputs.shape, 'f_h_t', f_h_t.shape, 'fps', fps.shape) 321 | for neuron in range(sample_batch_size): 322 | tps = fps[neuron*n_samples*batch_size:(neuron+1)*n_samples*batch_size] 323 | for img_i in range(batch_size): 324 | img_name = (labels[img_i + batch_size*input_i], img_i + batch_size*input_i) 325 | ps_key= (img_name, '{0}_{1}'.format(children[layer_i].__class__.__name__, layer_i), neuron+nt*sample_batch_size) 326 | ps = [tps[ img_i +batch_size*_] for _ in range(n_samples)] 327 | ps = np.asarray(ps) 328 | ps = ps.T 329 | # print('img i', img_i, input_i, batch_size, 'neuron', neuron, ps_key, ps.shape) 330 | all_ps[ps_key] = np.copy(ps) 331 | 332 | if input_i == 0 and nt == 0: 333 | os.system('nvidia-smi') 334 | del f_h_t_t 335 | del batch_data, inner_outputs 336 | torch.cuda.empty_cache() 337 | 338 | del temp_model1, temp_model2 339 | return all_ps, sample_layers 340 | 341 | 342 | def find_min_max(model_name, all_ps, sample_layers, cut_val=20, top_k = 10): 343 | max_ps = {} 344 | max_vals = [] 345 | n_classes = 0 346 | n_samples = 0 347 | for k in sorted(all_ps.keys()): 348 | #print(k) 349 | all_ps[k] = all_ps[k][:, :cut_val] 350 | n_classes = all_ps[k].shape[0] 351 | n_samples = all_ps[k].shape[1] 352 | # maximum increase diff 353 | 354 | vs = [] 355 | for l in range(num_classes): 356 | vs.append( np.amax(all_ps[k][l][1:]) - np.amin(all_ps[k][l][:1]) ) 357 | # vs.append( np.amax(all_ps[k][l]) - np.amin(all_ps[k][l]) ) 358 | # vs.append( np.amax(all_ps[k][l][all_ps[k].shape[1]//5:]) - np.amin(all_ps[k][l][:all_ps[k].shape[1]//5]) ) 359 | ml = np.argsort(np.asarray(vs))[-1] 360 | sml = np.argsort(np.asarray(vs))[-2] 361 | val = vs[ml] - vs[sml] 362 | # print(k, ml, sml) 363 | 364 | max_vals.append(val) 365 | max_ps[k] = (ml, val) 366 | 367 | neuron_ks = [] 368 | imgs = [] 369 | for k in sorted(max_ps.keys()): 370 | #print(k) 371 | nk = (k[1], k[2]) 372 | neuron_ks.append(nk) 373 | imgs.append(k[0]) 374 | #print(neuron_ks) 375 | neuron_ks = list(set(neuron_ks)) 376 | neuron_ks = sorted(neuron_ks,reverse=True) 377 | imgs = list(set(imgs)) 378 | #print('imgs', imgs) 379 | 380 | min_ps = {} 381 | min_vals = [] 382 | n_imgs = len(imgs) 383 | 384 | print(neuron_ks) 385 | 386 | for k in neuron_ks: 387 | vs = [] 388 | ls = [] 389 | vdict = {} 390 | for img in sorted(imgs): 391 | nk = (img, k[0], k[1]) 392 | l = max_ps[nk][0] 393 | v = max_ps[nk][1] 394 | vs.append(v) 395 | ls.append(l) 396 | if not ( l in vdict.keys() ): 397 | vdict[l] = [v] 398 | else: 399 | vdict[l].append(v) 400 | ml = max(set(ls), key=ls.count) 401 | 402 | 403 | fvs = [] 404 | # does not count when l not equal ml 405 | for img in sorted(imgs): 406 | img_l = int(img[0]) 407 | if img_l == ml: 408 | continue 409 | nk = (img, k[0], k[1]) 410 | l = max_ps[nk][0] 411 | v = max_ps[nk][1] 412 | if l != ml: 413 | continue 414 | fvs.append(v) 415 | 416 | # if ml == 2: 417 | # print(ml, k, len(fvs), np.amin(fvs), fvs) 418 | 419 | if len(fvs) > 0: 420 | min_ps[k] = (ml, ls.count(ml), np.amin(fvs), fvs) 421 | min_vals.append(np.amin(fvs)) 422 | # min_ps[k] = (ml, ls.count(ml), np.mean(fvs), fvs) 423 | # min_vals.append(np.mean(fvs)) 424 | 425 | else: 426 | min_ps[k] = (ml, 0, 0, fvs) 427 | min_vals.append(0) 428 | 429 | 430 | keys = min_ps.keys() 431 | keys = [] 432 | for k in min_ps.keys(): 433 | print(k) 434 | print(int(n_imgs * 0.9)) 435 | print(min_ps[k][1]) 436 | #if min_ps[k][1] >= int(n_imgs * 0.9): 437 | if arch == "vgg11bn": 438 | if min_ps[k][1] >= int(n_imgs * 0.7): 439 | keys.append(k) 440 | else: 441 | if min_ps[k][1] >= int(n_imgs * 0.9): 442 | keys.append(k) 443 | sorted_key = sorted(keys, key=lambda x: min_ps[x][2] ) 444 | 445 | print(sorted_key) 446 | 447 | if Print_Level > 0: 448 | print('n samples', n_samples, 'n class', n_classes, 'n_imgs', n_imgs) 449 | # print('sorted_key', sorted_key) 450 | 451 | 452 | neuron_dict = {} 453 | neuron_dict[model_name] = [] 454 | maxval = min_ps[sorted_key[-1]][2] 455 | layers = {} 456 | labels = {} 457 | allns = 0 458 | max_sampling_val = -np.inf 459 | 460 | # for i in range(len(sorted_key)): 461 | # k = sorted_key[-i-1] 462 | # layer = k[0] 463 | # neuron = k[1] 464 | # label = min_ps[k][0] 465 | # if layer not in layers.keys(): 466 | # layers[layer] = 1 467 | # else: 468 | # layers[layer] += 1 469 | # if layers[layer] <= 3: 470 | # if (layer, neuron, min_ps[k][0]) in neuron_dict[model_name]: 471 | # continue 472 | # if min_ps[k][2] > max_sampling_val: 473 | # max_sampling_val = min_ps[k][2] 474 | # if Print_Level > 0: 475 | # print('min max val across images', 'k', k, 'label', min_ps[k][0], min_ps[k][1], 'value', min_ps[k][2]) 476 | # if Print_Level > 1: 477 | # print(min_ps[k][3]) 478 | # allns += 1 479 | # neuron_dict[model_name].append( (layer, neuron, min_ps[k][0]) ) 480 | # if allns >= top_k: 481 | # # if allns > top_k//2: 482 | # break 483 | 484 | # early layers 485 | labels = {} 486 | for i in range(len(sorted_key)): 487 | k = sorted_key[-i-1] 488 | layer = k[0] 489 | neuron = k[1] 490 | label = min_ps[k][0] 491 | if (layer, neuron, min_ps[k][0]) in neuron_dict[model_name]: 492 | continue 493 | if label not in labels.keys(): 494 | labels[label] = 0 495 | if int(layer.split('_')[-1]) < sample_layers[-1] and labels[label] < 1: 496 | # if True: 497 | labels[label] += 1 498 | 499 | if min_ps[k][2] > max_sampling_val: 500 | max_sampling_val = min_ps[k][2] 501 | if Print_Level > 0: 502 | print('min max val across images', 'k', k, 'label', min_ps[k][0], min_ps[k][1], 'value', min_ps[k][2]) 503 | if Print_Level > 1: 504 | print(min_ps[k][3]) 505 | allns += 1 506 | neuron_dict[model_name].append( (layer, neuron, min_ps[k][0]) ) 507 | if allns >= top_k: 508 | break 509 | 510 | # last layers 511 | labels = {} 512 | for i in range(len(sorted_key)): 513 | k = sorted_key[-i-1] 514 | layer = k[0] 515 | neuron = k[1] 516 | label = min_ps[k][0] 517 | if (layer, neuron, min_ps[k][0]) in neuron_dict[model_name]: 518 | continue 519 | if label not in labels.keys(): 520 | labels[label] = 0 521 | if int(layer.split('_')[-1]) == sample_layers[-1] and labels[label] < 1: 522 | # if True: 523 | labels[label] += 1 524 | 525 | if min_ps[k][2] > max_sampling_val: 526 | max_sampling_val = min_ps[k][2] 527 | if Print_Level > 0: 528 | print('min max val across images', 'k', k, 'label', min_ps[k][0], min_ps[k][1], 'value', min_ps[k][2]) 529 | if Print_Level > 1: 530 | print(min_ps[k][3]) 531 | allns += 1 532 | neuron_dict[model_name].append( (layer, neuron, min_ps[k][0]) ) 533 | if allns >= top_k: 534 | break 535 | 536 | return neuron_dict, max_sampling_val 537 | 538 | def read_all_ps(model_name, all_ps, sample_layers, top_k=10, cut_val=20): 539 | return find_min_max(model_name, all_ps, sample_layers, cut_val, top_k=top_k) 540 | 541 | def filter_img(): 542 | mask = np.zeros((h, w), dtype=np.float32) 543 | Troj_w = int(np.sqrt(Troj_size) * 0.8) 544 | for i in range(h): 545 | for j in range(w): 546 | if j >= h/2 and j < h/2 + Troj_w \ 547 | and i >= w/2 and i < w/2 + Troj_w: 548 | mask[j,i] = 1 549 | return mask 550 | 551 | 552 | def nc_filter_img(): 553 | if use_mask: 554 | mask = np.zeros((h, w), dtype=np.float32) 555 | for i in range(h): 556 | for j in range(w): 557 | # if not( j >= w*1/4.0 and j < w*3/4.0 and i >= h*1/4.0 and i < h*3/4.0): 558 | if True: 559 | mask[i,j] = 1 560 | mask = np.zeros((h, w), dtype=np.float32) + 1 561 | else: 562 | mask = np.zeros((h, w), dtype=np.float32) + 1 563 | return mask 564 | 565 | 566 | def loss_fn(inner_outputs_b, inner_outputs_a, logits, con_mask, neuron, tlabel, acc, e, re_epochs): 567 | neuron_mask = torch.zeros([1, inner_outputs_a.shape[1],1,1]).cuda() 568 | neuron_mask[:,neuron,:,:] = 1 569 | vloss1 = torch.sum(inner_outputs_b * neuron_mask)/torch.sum(neuron_mask) 570 | vloss2 = torch.sum(inner_outputs_b * (1-neuron_mask))/torch.sum(1-neuron_mask) 571 | relu_loss1 = torch.sum(inner_outputs_a * neuron_mask)/torch.sum(neuron_mask) 572 | relu_loss2 = torch.sum(inner_outputs_a * (1-neuron_mask))/torch.sum(1-neuron_mask) 573 | 574 | vloss3 = torch.sum(inner_outputs_b * torch.lt(inner_outputs_b, 0) )/torch.sum(1-neuron_mask) 575 | 576 | loss = - vloss1 - relu_loss1 + 0.0001 * vloss2 + 0.0001 * relu_loss2 577 | #print(vloss1,relu_loss1,vloss2,relu_loss2) 578 | # if vloss1 < 0 and e < try_epochs: 579 | # print('use vloss3', e) 580 | # loss = - vloss3# - vloss1 581 | # else: 582 | # print('not use vloss3', e) 583 | mask_loss = torch.sum(con_mask) 584 | mask_nz = torch.sum(torch.gt(con_mask, mask_epsilon)) 585 | mask_cond1 = torch.gt(mask_nz, Troj_size) 586 | mask_cond2 = torch.gt(mask_nz, Troj_size * 1.2) 587 | # mask_cond1 = torch.gt(mask_loss, Troj_size) 588 | # mask_cond2 = torch.gt(mask_loss, Troj_size * 1.2) 589 | mask_add_loss = torch.where(mask_cond1, torch.where(mask_cond2, 2 * re_mask_weight * mask_loss, 1 * re_mask_weight * mask_loss), 0.01 * mask_loss) 590 | loss += mask_add_loss 591 | logits_loss = torch.sum(logits[:,tlabel]) 592 | loss += - 1 * logits_loss 593 | 594 | print("vloss1",vloss1,"\nrelu_loss1",relu_loss1,"\nvloss2",vloss2,"\nrelu_loss2",relu_loss2,"\nmask_loss",mask_loss,"\nmask_add_loss",mask_add_loss,"\nlogits_loss",logits_loss,"\n") 595 | 596 | if e > re_epochs//2: 597 | # if True: 598 | loss += - 2 * logits_loss 599 | return loss, vloss1, vloss2, vloss3, relu_loss1, relu_loss2, mask_loss, mask_nz, mask_add_loss, logits_loss 600 | 601 | def reverse_engineer(model_type, model, children, oimages, olabels, weights_file, Troj_Layer, Troj_Neuron, Troj_Label, Troj_size, re_epochs): 602 | 603 | before_block = [] 604 | def get_before_block(): 605 | def hook(model, input, output): 606 | for ip in input: 607 | before_block.append( ip.clone() ) 608 | return hook 609 | 610 | after_bn3 = [] 611 | def get_after_bn3(): 612 | def hook(model, input, output): 613 | for ip in output: 614 | after_bn3.append( ip.clone() ) 615 | return hook 616 | 617 | after_iden = [] 618 | def get_after_iden(): 619 | def hook(model, input, output): 620 | for ip in output: 621 | after_iden.append( ip.clone() ) 622 | return hook 623 | 624 | after_bns = [] 625 | def get_after_bns(): 626 | def hook(model, input, output): 627 | for ip in output: 628 | after_bns.append( ip.clone() ) 629 | return hook 630 | 631 | 632 | re_batch_size = config['re_batch_size'] 633 | if model_type == 'DenseNet': 634 | re_batch_size = max(re_batch_size // 3, 1) 635 | if model_type == 'ResNet': 636 | re_batch_size = max(re_batch_size // 2, 1) 637 | if re_batch_size > len(oimages): 638 | re_batch_size = len(oimages) 639 | 640 | handles = [] 641 | if model_type == 'ResNet': 642 | children_modules = list(list(children[Troj_Layer].children())[-1].children()) 643 | print(len(children_modules), children_modules) 644 | bn3_module = children_modules[4] 645 | handle = bn3_module.register_forward_hook(get_after_bn3()) 646 | handles.append(handle) 647 | if len(children_modules) > 5: 648 | iden_module = children_modules[-1] 649 | handle = iden_module.register_forward_hook(get_after_iden()) 650 | handles.append(handle) 651 | else: 652 | iden_module = children_modules[0] 653 | handle = iden_module.register_forward_hook(get_before_block()) 654 | handles.append(handle) 655 | elif model_type == 'Inception3': 656 | children_modules = [] 657 | for j in range(len(list(children[Troj_Layer].modules()))): 658 | tm = list(children[Troj_Layer].modules())[j] 659 | if len(list(tm.children())) == 0: 660 | children_modules.append(tm) 661 | for j in range(len(children_modules)): 662 | print(j, children_modules[j]) 663 | if children[Troj_Layer].__class__.__name__ == 'InceptionA': 664 | tmodule1 = children_modules[1] 665 | tmodule2 = children_modules[5] 666 | tmodule3 = children_modules[11] 667 | tmodule4 = children_modules[13] 668 | handle = tmodule1.register_forward_hook(get_after_bns()) 669 | handles.append(handle) 670 | handle = tmodule2.register_forward_hook(get_after_bns()) 671 | handles.append(handle) 672 | handle = tmodule3.register_forward_hook(get_after_bns()) 673 | handles.append(handle) 674 | handle = tmodule4.register_forward_hook(get_after_bns()) 675 | handles.append(handle) 676 | elif children[Troj_Layer].__class__.__name__ == 'InceptionB': 677 | tmodule1 = children_modules[1] 678 | tmodule2 = children_modules[7] 679 | tmodule3 = children_modules[0] 680 | handle = tmodule1.register_forward_hook(get_after_bns()) 681 | handles.append(handle) 682 | handle = tmodule2.register_forward_hook(get_after_bns()) 683 | handles.append(handle) 684 | handle = tmodule3.register_forward_hook(get_before_block()) 685 | handles.append(handle) 686 | elif children[Troj_Layer].__class__.__name__ == 'InceptionC': 687 | tmodule1 = children_modules[1] 688 | tmodule2 = children_modules[7] 689 | tmodule3 = children_modules[17] 690 | tmodule4 = children_modules[19] 691 | handle = tmodule1.register_forward_hook(get_after_bns()) 692 | handles.append(handle) 693 | handle = tmodule2.register_forward_hook(get_after_bns()) 694 | handles.append(handle) 695 | handle = tmodule3.register_forward_hook(get_after_bns()) 696 | handles.append(handle) 697 | handle = tmodule4.register_forward_hook(get_after_bns()) 698 | handles.append(handle) 699 | elif children[Troj_Layer].__class__.__name__ == 'InceptionD': 700 | tmodule1 = children_modules[3] 701 | tmodule2 = children_modules[11] 702 | tmodule3 = children_modules[0] 703 | handle = tmodule1.register_forward_hook(get_after_bns()) 704 | handles.append(handle) 705 | handle = tmodule2.register_forward_hook(get_after_bns()) 706 | handles.append(handle) 707 | handle = tmodule3.register_forward_hook(get_before_block()) 708 | handles.append(handle) 709 | elif children[Troj_Layer].__class__.__name__ == 'InceptionE': 710 | tmodule1 = children_modules[1] 711 | tmodule2 = children_modules[5] 712 | tmodule3 = children_modules[7] 713 | tmodule4 = children_modules[13] 714 | tmodule5 = children_modules[15] 715 | tmodule6 = children_modules[17] 716 | handle = tmodule1.register_forward_hook(get_after_bns()) 717 | handles.append(handle) 718 | handle = tmodule2.register_forward_hook(get_after_bns()) 719 | handles.append(handle) 720 | handle = tmodule3.register_forward_hook(get_after_bns()) 721 | handles.append(handle) 722 | handle = tmodule4.register_forward_hook(get_after_bns()) 723 | handles.append(handle) 724 | handle = tmodule5.register_forward_hook(get_after_bns()) 725 | handles.append(handle) 726 | handle = tmodule6.register_forward_hook(get_after_bns()) 727 | handles.append(handle) 728 | elif model.__class__.__name__ == 'DenseNet': 729 | target_module = list(children[Troj_Layer].modules())[-1] 730 | handle = target_module.register_forward_hook(get_after_bns()) 731 | handles.append(handle) 732 | elif model.__class__.__name__ == 'Net': 733 | target_module = list(children[Troj_Layer].modules())[-1] 734 | #print(target_module) 735 | handle = target_module.register_forward_hook(get_after_bns()) 736 | handles.append(handle) 737 | elif model.__class__.__name__ == 'VGG': 738 | target_module = list(children[Troj_Layer].modules())[-1] 739 | #print(target_module) 740 | handle = target_module.register_forward_hook(get_after_bns()) 741 | handles.append(handle) 742 | elif model.__class__.__name__ == 'LeNet5': 743 | target_module = list(children[Troj_Layer].modules())[-1] 744 | #print(target_module) 745 | handle = target_module.register_forward_hook(get_after_bns()) 746 | handles.append(handle) 747 | elif model.__class__.__name__ == 'AlexNet': 748 | target_module = list(children[Troj_Layer].modules())[-1] 749 | #print(target_module) 750 | handle = target_module.register_forward_hook(get_after_bns()) 751 | handles.append(handle) 752 | print('Target Layer', Troj_Layer, children[Troj_Layer], 'Neuron', Troj_Neuron, 'Target Label', Troj_Label) 753 | 754 | # delta = torch.randn(1,3,h,w).cuda() 755 | delta = torch.rand(1,3,h,w).cuda() * 2 - 1 756 | print(torch.rand(1)) 757 | print(torch.rand(1)) 758 | print(delta) 759 | mask = filter_img().reshape((1,1,h,w)) * 6 - 4 760 | mask= torch.FloatTensor(mask).cuda() 761 | delta.requires_grad = True 762 | mask.requires_grad = True 763 | optimizer = torch.optim.Adam([delta, mask], lr=re_mask_lr) 764 | print('before optimizing',) 765 | for e in range(re_epochs): 766 | print("reverse_engineer e:",e) 767 | facc = 0 768 | flogits = [] 769 | p = np.random.permutation(oimages.shape[0]) 770 | images = oimages[p] 771 | labels = olabels[p] 772 | #print(len(images)//re_batch_size) 773 | for i in range(len(images)//re_batch_size): 774 | optimizer.zero_grad() 775 | model.zero_grad() 776 | after_bn3.clear() 777 | before_block.clear() 778 | after_iden.clear() 779 | after_bns.clear() 780 | 781 | batch_data = torch.FloatTensor(images[re_batch_size*i:re_batch_size*(i+1)]) 782 | batch_data = batch_data.cuda() 783 | 784 | con_mask = torch.tanh(mask)/2.0 + 0.5 785 | # delta is bounded by 0 and 1 786 | # use_delta = torch.clamp(delta, min=0.0 , max=1.0) 787 | # if e < try_epochs: 788 | # use_delta = torch.tanh(torch.clamp(delta, -1, 1))/2.0 + 0.5 789 | # else: 790 | # use_delta = torch.tanh(delta)/2.0 + 0.5 791 | use_delta = torch.tanh(delta)/2.0 + 0.5 792 | use_mask = con_mask 793 | in_data = use_mask * use_delta + (1-use_mask) * batch_data 794 | 795 | if model_type == 'ResNet': 796 | 797 | logits = model(in_data) 798 | logits_np = logits.cpu().detach().numpy() 799 | 800 | after_bn3_t = torch.stack(after_bn3, 0) 801 | iden = None 802 | if len(before_block) > 0: 803 | iden = before_block[0] 804 | else: 805 | after_iden_t = torch.stack(after_iden, 0) 806 | iden = after_iden_t 807 | inner_outputs_b = iden + after_bn3_t 808 | inner_outputs_a = F.relu(inner_outputs_b) 809 | 810 | elif model_type == 'Inception3': 811 | 812 | logits = model(in_data) 813 | logits_np = logits.cpu().detach().numpy() 814 | 815 | if children[Troj_Layer].__class__.__name__ == 'InceptionA': 816 | after_bn1_t = torch.stack(after_bns[0*re_batch_size:1*re_batch_size], 0) 817 | after_bn2_t = torch.stack(after_bns[1*re_batch_size:2*re_batch_size], 0) 818 | after_bn3_t = torch.stack(after_bns[2*re_batch_size:3*re_batch_size], 0) 819 | after_bn4_t = torch.stack(after_bns[3*re_batch_size:4*re_batch_size], 0) 820 | inner_outputs_b = torch.cat([after_bn1_t, after_bn2_t, after_bn3_t, after_bn4_t], 1) 821 | if children[Troj_Layer].__class__.__name__ == 'InceptionB': 822 | after_bn1_t = torch.stack(after_bns[0*re_batch_size:1*re_batch_size], 0) 823 | after_bn2_t = torch.stack(after_bns[1*re_batch_size:2*re_batch_size], 0) 824 | before_in_t = before_block[0] 825 | branch_pool = F.max_pool2d(before_in_t, kernel_size=3, stride=2) 826 | inner_outputs_b = torch.cat([after_bn1_t, after_bn2_t, branch_pool], 1) 827 | if children[Troj_Layer].__class__.__name__ == 'InceptionC': 828 | after_bn1_t = torch.stack(after_bns[0*re_batch_size:1*re_batch_size], 0) 829 | after_bn2_t = torch.stack(after_bns[1*re_batch_size:2*re_batch_size], 0) 830 | after_bn3_t = torch.stack(after_bns[2*re_batch_size:3*re_batch_size], 0) 831 | after_bn4_t = torch.stack(after_bns[3*re_batch_size:4*re_batch_size], 0) 832 | inner_outputs_b = torch.cat([after_bn1_t, after_bn2_t, after_bn3_t, after_bn4_t], 1) 833 | if children[Troj_Layer].__class__.__name__ == 'InceptionD': 834 | after_bn1_t = torch.stack(after_bns[0*re_batch_size:1*re_batch_size], 0) 835 | after_bn2_t = torch.stack(after_bns[1*re_batch_size:2*re_batch_size], 0) 836 | before_in_t = before_block[0] 837 | branch_pool = F.max_pool2d(before_in_t, kernel_size=3, stride=2) 838 | inner_outputs_b = torch.cat([after_bn1_t, after_bn2_t, branch_pool], 1) 839 | if children[Troj_Layer].__class__.__name__ == 'InceptionE': 840 | after_bn1_t = torch.stack(after_bns[0*re_batch_size:1*re_batch_size], 0) 841 | after_bn2_t = torch.stack(after_bns[1*re_batch_size:2*re_batch_size], 0) 842 | after_bn3_t = torch.stack(after_bns[2*re_batch_size:3*re_batch_size], 0) 843 | after_bn4_t = torch.stack(after_bns[3*re_batch_size:4*re_batch_size], 0) 844 | after_bn5_t = torch.stack(after_bns[4*re_batch_size:5*re_batch_size], 0) 845 | after_bn6_t = torch.stack(after_bns[5*re_batch_size:6*re_batch_size], 0) 846 | inner_outputs_b = torch.cat([after_bn1_t, after_bn2_t, after_bn3_t, after_bn4_t, after_bn5_t, after_bn6_t], 1) 847 | 848 | inner_outputs_a = F.relu(inner_outputs_b) 849 | 850 | elif model_type == 'DenseNet': 851 | 852 | logits = model(in_data) 853 | logits_np = logits.cpu().detach().numpy() 854 | inner_outputs_b = torch.stack(after_bns, 0) 855 | inner_outputs_a = F.relu(inner_outputs_b) 856 | elif model_type == 'Net': 857 | 858 | logits = model(in_data) 859 | #print(F.softmax(logits)) 860 | #print(F.softmax(logits).max(1)) 861 | logits_np = logits.cpu().detach().numpy() 862 | inner_outputs_b = torch.stack(after_bns, 0) 863 | inner_outputs_a = F.relu(inner_outputs_b) 864 | elif model_type == 'VGG': 865 | 866 | logits = model(in_data) 867 | #print(F.softmax(logits)) 868 | #print(F.softmax(logits).max(1)) 869 | logits_np = logits.cpu().detach().numpy() 870 | inner_outputs_b = torch.stack(after_bns, 0) 871 | inner_outputs_a = F.relu(inner_outputs_b) 872 | 873 | elif model_type == 'LeNet5': 874 | 875 | logits = model(in_data) 876 | #print(F.softmax(logits)) 877 | #print(F.softmax(logits).max(1)) 878 | logits_np = logits.cpu().detach().numpy() 879 | inner_outputs_b = torch.stack(after_bns, 0) 880 | inner_outputs_a = F.relu(inner_outputs_b) 881 | 882 | elif model_type == 'AlexNet': 883 | 884 | logits = model(in_data) 885 | #print(F.softmax(logits)) 886 | #print(F.softmax(logits).max(1)) 887 | logits_np = logits.cpu().detach().numpy() 888 | inner_outputs_b = torch.stack(after_bns, 0) 889 | inner_outputs_a = F.relu(inner_outputs_b) 890 | 891 | flogits.append(logits_np) 892 | loss, vloss1, vloss2, vloss3, relu_loss1, relu_loss2, mask_loss, mask_nz, mask_add_loss, logits_loss\ 893 | = loss_fn(inner_outputs_b, inner_outputs_a, logits, use_mask, Troj_Neuron, int(Troj_Label), facc, e, re_epochs) 894 | if e > 0: 895 | loss.backward(retain_graph=True) 896 | optimizer.step() 897 | print(loss) 898 | # break 899 | flogits = np.concatenate(flogits, axis=0) 900 | preds = np.argmax(flogits, axis=1) 901 | 902 | # do not change Troj_Label 903 | Troj_Label2 = np.argmax(np.bincount(preds)) 904 | 905 | facc = np.sum(preds == Troj_Label2) / float(preds.shape[0]) 906 | 907 | '''if e % 10 == 0: 908 | print(e, 'loss', loss.cpu().detach().numpy(), 'acc {:.4f}'.format(facc),'target label', int(Troj_Label), int(Troj_Label2), 'logits_loss', logits_loss.cpu().detach().numpy(),\ 909 | 'vloss1', vloss1.cpu().detach().numpy(), 'vloss2', vloss2.cpu().detach().numpy(),\ 910 | 'relu_loss1', relu_loss1.cpu().detach().numpy(), 'max relu_loss1', np.amax(inner_outputs_a.cpu().detach().numpy()),\ 911 | 'relu_loss2', relu_loss2.cpu().detach().numpy(),\ 912 | 'mask_loss', mask_loss.cpu().detach().numpy(), 'mask_nz', mask_nz.cpu().detach().numpy(), 'mask_add_loss', mask_add_loss.cpu().detach().numpy()) 913 | print('labels', flogits[:5,:]) 914 | print('logits', np.argmax(flogits, axis=1)) 915 | print('delta', use_delta[0,0,:5,:5]) 916 | print('mask', use_mask[0,0,:5,:5])''' 917 | 918 | if facc > 0.99: 919 | break 920 | delta = use_delta.cpu().detach().numpy() 921 | con_mask = use_mask.cpu().detach().numpy() 922 | adv = in_data.cpu().detach().numpy() 923 | adv = deprocess(adv) 924 | 925 | #os.system('nvidia-smi') 926 | # cleaning up 927 | for handle in handles: 928 | handle.remove() 929 | 930 | return facc, adv, delta, con_mask, Troj_Label2 931 | 932 | def re_mask(model_type, model, neuron_dict, children, images, labels, scratch_dirpath, re_epochs): 933 | validated_results = [] 934 | for key in sorted(neuron_dict.keys()): 935 | weights_file = key 936 | for task in neuron_dict[key]: 937 | Troj_Layer, Troj_Neuron, samp_label = task 938 | Troj_Neuron = int(Troj_Neuron) 939 | Troj_Layer = int(Troj_Layer.split('_')[1]) 940 | 941 | print(weights_file) 942 | RE_img = os.path.join(scratch_dirpath,'imgs', '{0}_model_{1}_{2}_{3}_{4}.png'.format( weights_file.split('/')[-2], Troj_Layer, Troj_Neuron, Troj_size, samp_label)) 943 | RE_mask = os.path.join(scratch_dirpath,'masks', '{0}_model_{1}_{2}_{3}_{4}.pkl'.format( weights_file.split('/')[-2], Troj_Layer, Troj_Neuron, Troj_size, samp_label)) 944 | RE_delta = os.path.join(scratch_dirpath,'deltas', '{0}_model_{1}_{2}_{3}_{4}.pkl'.format(weights_file.split('/')[-2], Troj_Layer, Troj_Neuron, Troj_size, samp_label)) 945 | 946 | max_acc = 0 947 | max_results = [] 948 | for i in range(mask_multi_start): 949 | acc, rimg, rdelta, rmask, optz_label = reverse_engineer(model_type, model, children, images, labels, weights_file, Troj_Layer, Troj_Neuron, samp_label, Troj_size, re_epochs) 950 | 951 | #print(rimg.shape) 952 | re_save_path = os.path.join(scratch_dirpath,'imgs', '{0}_model_{1}_{2}_{3}_{4}_{5}_{6}.png'.format( weights_file.split('/')[-2], Troj_Layer, Troj_Neuron, Troj_size, samp_label, optz_label, acc)) 953 | skimage.io.imsave(re_save_path, rimg[0]) 954 | 955 | #skimage.io.imsave(RE_img, rimg[0]) 956 | #skimage.io.imsave(RE_mask, rmask.transpose(0,2,3,1)[0]) 957 | #print(rdelta.shape) 958 | #skimage.io.imsave(RE_delta, rdelta.transpose(0,2,3,1)[0]) 959 | # clear cache 960 | torch.cuda.empty_cache() 961 | 962 | if Print_Level > 0: 963 | print('RE mask', Troj_Layer, Troj_Neuron, 'Label', optz_label, 'RE acc', acc) 964 | if acc > max_acc: 965 | max_acc = acc 966 | max_results = (rimg, rdelta, rmask, optz_label, RE_img, RE_mask, RE_delta, samp_label, acc) 967 | if max_acc >= reasr_bound - 0.2: 968 | validated_results.append( max_results ) 969 | if max_acc > 0.99 and optz_label == samp_label: 970 | break 971 | return validated_results 972 | 973 | 974 | def stamp(n_img, delta, mask): 975 | mask0 = nc_filter_img() 976 | mask = mask * mask0 977 | r_img = n_img.copy() 978 | mask = mask.reshape((1,1,h,w)) 979 | r_img = n_img * (1-mask) + delta * mask 980 | #print("stamp\n\n\n") 981 | #print(delta) 982 | #print(n_img) 983 | return r_img 984 | 985 | def filter_stamp(n_img, trigger): 986 | if channel_last: 987 | t_image = tf.placeholder(tf.float32, shape=(None, h, w, 3)) 988 | ti_image = t_image 989 | else: 990 | t_image = tf.placeholder(tf.float32, shape=(None, 3, h, w)) 991 | ti_image = tf.transpose(t_image, [0,2,3,1]) 992 | 993 | tdelta = tf.placeholder(tf.float32, shape=(12, 3)) 994 | imax = tf.nn.max_pool( ti_image, ksize=[1,window_size,window_size,1], strides=[1,1,1,1], padding='SAME') 995 | imin = -tf.nn.max_pool(-ti_image, ksize=[1,window_size,window_size,1], strides=[1,1,1,1], padding='SAME') 996 | iavg = tf.nn.avg_pool( ti_image, ksize=[1,window_size,window_size,1], strides=[1,1,1,1], padding='SAME') 997 | i_image = tf.reshape( tf.matmul( tf.reshape( tf.concat([ti_image, imax, imin, iavg], axis=3), (-1,12)) , tdelta), [-1,h,w,3]) 998 | 999 | if not channel_last: 1000 | i_image = tf.transpose(i_image, [0,3,1,2]) 1001 | 1002 | with tf.Session() as sess: 1003 | r_img = sess.run(i_image, {t_image: n_img, tdelta:trigger}) 1004 | return r_img 1005 | 1006 | def test(model, model_type, test_xs, test_ys, result, scratch_dirpath, test_index, log_dir, log_name, mode='mask'): 1007 | 1008 | re_batch_size = config['re_batch_size'] 1009 | if model_type == 'DenseNet': 1010 | re_batch_size = max(re_batch_size // 3, 1) 1011 | if model_type == 'ResNet': 1012 | re_batch_size = max(re_batch_size // 2, 1) 1013 | if re_batch_size > len(test_xs): 1014 | re_batch_size = len(test_xs) 1015 | 1016 | clean_images = test_xs 1017 | 1018 | # if mode == 'mask': 1019 | if True: 1020 | rimg, rdelta, rmask, tlabel = result[:4] 1021 | rmask = rmask * rmask > mask_epsilon 1022 | t_images = stamp(clean_images, rdelta, rmask) 1023 | 1024 | saved_images = deprocess(t_images) 1025 | 1026 | rt_images = t_images 1027 | 1028 | if Print_Level > 0: 1029 | print(np.amin(rt_images), np.amax(rt_images)) 1030 | 1031 | yt = np.zeros(len(rt_images)).astype(np.int32) + tlabel 1032 | fpreds = [] 1033 | for i in range(len(rt_images)//re_batch_size): 1034 | batch_data = torch.FloatTensor(rt_images[re_batch_size*i:re_batch_size*(i+1)]) 1035 | batch_data = batch_data.cuda() 1036 | preds = model(batch_data) 1037 | fpreds.append(preds.cpu().detach().numpy()) 1038 | fpreds = np.concatenate(fpreds) 1039 | 1040 | preds = np.argmax(fpreds, axis=1) 1041 | #print(preds) 1042 | #print(test_ys) 1043 | 1044 | if score_type == "mean": 1045 | print("tlabel:",tlabel) 1046 | for i in range(num_classes): 1047 | print("class:",i) 1048 | #print(preds) 1049 | #print(test_ys) 1050 | class_i_index = np.where(test_ys==i) 1051 | pred_class_i = preds[class_i_index] 1052 | score = float(np.sum(tlabel == pred_class_i))/float(pred_class_i.shape[0]) 1053 | print(score) 1054 | score = float(np.sum(tlabel == preds))/float(yt.shape[0]) 1055 | 1056 | if score_type == "mean_exclude_target": 1057 | print("tlabel:",tlabel) 1058 | correct_num = 0 1059 | total_num = 0 1060 | for i in range(num_classes): 1061 | print("class:",i) 1062 | print(len(preds)) 1063 | print(len(test_ys)) 1064 | class_i_index = np.where(test_ys==i) 1065 | print(class_i_index) 1066 | pred_class_i = preds[class_i_index] 1067 | score = float(np.sum(tlabel == pred_class_i))/float(pred_class_i.shape[0]) 1068 | if i != tlabel: 1069 | correct_num = correct_num + float(np.sum(tlabel == pred_class_i)) 1070 | total_num = total_num + float(pred_class_i.shape[0]) 1071 | 1072 | print(score) 1073 | 1074 | score = float(correct_num)/float(total_num) 1075 | #score = float(np.sum(tlabel == preds))/float(yt.shape[0]) 1076 | 1077 | elif score_type == "max": 1078 | print("tlabel:",tlabel) 1079 | score_list = [] 1080 | for i in range(num_classes): 1081 | print("class:",i) 1082 | #print(preds) 1083 | #print(test_ys) 1084 | class_i_index = np.where(test_ys==i) 1085 | pred_class_i = preds[class_i_index] 1086 | score = float(np.sum(tlabel == pred_class_i))/float(pred_class_i.shape[0]) 1087 | print(score) 1088 | if i != tlabel: 1089 | score_list.append(score) 1090 | 1091 | #score = float(np.sum(tlabel == preds))/float(yt.shape[0]) 1092 | score = max(score_list) 1093 | 1094 | 1095 | 1096 | #print('label', tlabel, 'score', score) 1097 | with open( log_dir+log_name ,'a') as f: 1098 | f.write('index:' + str(test_index)+' label:'+str(tlabel)+' score:'+str(score)) 1099 | f.write('\n') 1100 | test_result = str(tlabel)+"_"+str(score) 1101 | 1102 | return score,test_result 1103 | 1104 | def run_abs(model_filepath, result_filepath, scratch_dirpath, examples_dirpath, log_dir, log_name, arch_local="nin", example_img_format='png', troj_size_override=None, num_classes_override = None, total_size_override = None, re_bs_override = None, dpsgd_convert=0 ): 1105 | start = time.time() 1106 | 1107 | seed_torch(333) 1108 | 1109 | print(torch.rand(1)) 1110 | print(torch.rand(1)) 1111 | 1112 | global arch 1113 | arch = arch_local 1114 | 1115 | if troj_size_override: 1116 | global Troj_size 1117 | Troj_size = troj_size_override 1118 | 1119 | if num_classes_override: 1120 | global num_classes 1121 | num_classes = num_classes_override 1122 | 1123 | if total_size_override: 1124 | global w 1125 | global h 1126 | w = total_size_override 1127 | h = total_size_override 1128 | 1129 | if re_bs_override: 1130 | global config 1131 | config['re_batch_size'] = re_bs_override 1132 | 1133 | 1134 | 1135 | # create dirs 1136 | os.system('mkdir -p {0}'.format(os.path.join(scratch_dirpath, 'imgs'))) 1137 | os.system('mkdir -p {0}'.format(os.path.join(scratch_dirpath, 'masks'))) 1138 | os.system('mkdir -p {0}'.format(os.path.join(scratch_dirpath, 'temps'))) 1139 | os.system('mkdir -p {0}'.format(os.path.join(scratch_dirpath, 'deltas'))) 1140 | 1141 | fns = [os.path.join(examples_dirpath, fn) for fn in os.listdir(examples_dirpath) if fn.endswith(example_img_format)] 1142 | #random.shuffle(fns) 1143 | imgs = [] 1144 | fys = [] 1145 | image_mins = [] 1146 | image_maxs = [] 1147 | #print(fns) 1148 | mean = [0.4914, 0.4822, 0.4465] 1149 | std = [0.2023, 0.1994, 0.2010] 1150 | for fn in fns: 1151 | # read the image (using skimage) 1152 | img = skimage.io.imread(fn) 1153 | fys.append(int(fn.split('_')[-3])) 1154 | # perform tensor formatting and normalization explicitly 1155 | # convert to CHW dimension ordering 1156 | img = np.transpose(img, (2, 0, 1)) 1157 | 1158 | img = img/255 1159 | 1160 | for i in range(len(mean)): 1161 | img[i] = (img[i] - mean[i])/std[i] 1162 | 1163 | img = np.expand_dims(img, 0) 1164 | 1165 | image_mins.append(np.min(img)) 1166 | image_maxs.append(np.max(img)) 1167 | imgs.append(img) 1168 | fxs = np.concatenate(imgs) 1169 | fys = np.array(fys) 1170 | image_min = np.mean(image_mins) 1171 | image_max = np.mean(image_maxs) 1172 | 1173 | print('number of seed images', len(fys), fys.shape) 1174 | 1175 | test_xs = fxs 1176 | test_ys = fys 1177 | 1178 | n_sample_imgs = 200 1179 | sample_xs = [] 1180 | sample_ys = [] 1181 | sample_slots = np.zeros(num_classes) 1182 | for i in range(len(fys)): 1183 | #print(fys[i]) 1184 | if sample_slots[fys[i]] < n_sample_imgs//num_classes: 1185 | sample_xs.append(fxs[i]) 1186 | sample_ys.append(fys[i]) 1187 | sample_slots[fys[i]] += 1 1188 | if np.sum(sample_slots) >= n_sample_imgs: 1189 | break 1190 | sample_xs = np.array(sample_xs) 1191 | sample_ys = np.array(sample_ys) 1192 | 1193 | xs = sample_xs.copy() 1194 | ys = sample_ys.copy() 1195 | optz_xs = [] 1196 | optz_ys = [] 1197 | optz_slots = np.zeros(num_classes) 1198 | for i in range(len(fys)): 1199 | if optz_slots[fys[i]] < n_re_samples//num_classes: 1200 | optz_xs.append(fxs[i]) 1201 | optz_ys.append(fys[i]) 1202 | optz_slots[fys[i]] += 1 1203 | if np.sum(optz_slots) >= n_re_samples: 1204 | break 1205 | optz_xs = np.array(optz_xs) 1206 | optz_ys = np.array(optz_ys) 1207 | xs = optz_xs 1208 | ys = optz_ys 1209 | 1210 | if Print_Level > 0: 1211 | print('# samples for RE', len(ys)) 1212 | 1213 | if arch == "nin": 1214 | model = Net() 1215 | elif arch == "vgg16": 1216 | model = vgg.vgg16(num_classes=num_classes) 1217 | elif arch == "vgg11": 1218 | model = vgg.vgg11(num_classes=num_classes) 1219 | elif arch == "vgg11bn": 1220 | model = vgg_trojai.vgg11_bn(num_classes=num_classes) 1221 | elif arch == "resnet18": 1222 | model = resnet.resnet18(num_classes=num_classes) 1223 | 1224 | try: 1225 | #print(1) 1226 | model.load_state_dict(torch.load(model_filepath)) 1227 | #model.load_state_dict(torch.load(model_filepath),strict=True) 1228 | except: 1229 | model = torch.load(model_filepath) 1230 | #print(model) 1231 | 1232 | 1233 | 1234 | model = model.cuda() 1235 | 1236 | model.eval() 1237 | correct = 0 1238 | for i in range(len(test_xs)//10): 1239 | batch_data = torch.FloatTensor(test_xs[10*i:10*(i+1)]) 1240 | batch_data = batch_data.cuda() 1241 | 1242 | preds = model(batch_data) 1243 | 1244 | labels = torch.Tensor(test_ys[10*i:10*(i+1)]).cuda() 1245 | labels = labels.view(10,-1) 1246 | 1247 | preds = preds.data.max(1, keepdim=True)[1] 1248 | 1249 | correct += preds.eq(labels).cpu().sum() 1250 | 1251 | print(correct) 1252 | print(len(test_xs)) 1253 | 1254 | target_layers = [] 1255 | model_type = model.__class__.__name__ 1256 | if model_type == 'ResNet': 1257 | children = list(model.children()) 1258 | 1259 | children.insert(-1, torch.nn.Flatten()) 1260 | #print(children) 1261 | # target_layers = ['Bottleneck'] 1262 | target_layers = ['Sequential'] 1263 | #target_layers = ['AvgPool2d'] 1264 | elif model_type == 'Inception3': 1265 | children = list(model.children()) 1266 | children.insert(3, torch.nn.MaxPool2d(kernel_size=3, stride=2)) 1267 | children.insert(6, torch.nn.MaxPool2d(kernel_size=3, stride=2)) 1268 | children.insert(-1, torch.nn.AdaptiveAvgPool2d((1, 1))) 1269 | children.insert(-1, torch.nn.Flatten()) 1270 | target_layers = ['InceptionA', 'InceptionB', 'InceptionC', 'InceptionD', 'InceptionE'] 1271 | elif model_type == 'DenseNet': 1272 | children = list(model.children()) 1273 | children = list(children[0].children()) + children[1:] 1274 | nchildren = [] 1275 | for c in children: 1276 | if c.__class__.__name__ == '_Transition': 1277 | nchildren += list(c.children()) 1278 | else: 1279 | nchildren.append(c) 1280 | children = nchildren 1281 | children.insert(-1, torch.nn.ReLU(inplace=True)) 1282 | children.insert(-1, torch.nn.AdaptiveAvgPool2d((1, 1))) 1283 | children.insert(-1, torch.nn.Flatten()) 1284 | print(children) 1285 | target_layers = ['BatchNorm2d'] 1286 | elif model_type == 'Net': 1287 | children = list(model.children()) 1288 | children = list(children[0].children()) + children[1:] 1289 | 1290 | children.append(torch.nn.Flatten()) 1291 | print(children) 1292 | target_layers = ['Conv2d'] 1293 | elif model_type == 'VGG': 1294 | children = list(model.children()) 1295 | children = list(children[0].children()) + children[1:] 1296 | 1297 | children.insert(-1, torch.nn.Flatten()) 1298 | print(children) 1299 | target_layers = ['Conv2d'] 1300 | 1301 | elif model_type == 'LeNet5': 1302 | children = list(model.children()) 1303 | 1304 | children.insert(-3, torch.nn.Flatten()) 1305 | print(children) 1306 | target_layers = ['Conv2d'] 1307 | elif model_type == 'AlexNet': 1308 | #children = list(model.children()) 1309 | children = list(model.children()) 1310 | children = list(children[0].children()) + list(children[1].children()) 1311 | 1312 | children.insert(-7, torch.nn.Flatten()) 1313 | print(children) 1314 | target_layers = ['Conv2d'] 1315 | else: 1316 | print('other model', model_type) 1317 | sys.exit() 1318 | 1319 | if Print_Level > 0: 1320 | print('layers') 1321 | for i in range(len(children)): 1322 | print(i, children[i], type(children[i])) 1323 | 1324 | if Print_Level > 0: 1325 | print('image range', np.amin(test_xs), np.amax(test_xs)) 1326 | 1327 | neuron_dict = {} 1328 | sampling_val = 0 1329 | 1330 | maxes = check_values( sample_xs, sample_ys, model, children, target_layers) 1331 | torch.cuda.empty_cache() 1332 | all_ps, sample_layers = sample_neuron(sample_xs, sample_ys, model, children, target_layers, model_type, maxes) 1333 | torch.cuda.empty_cache() 1334 | 1335 | neuron_dict, sampling_val = read_all_ps(model_filepath, all_ps, sample_layers, top_k = top_n_neurons) 1336 | print('Compromised Neuron Candidates (Layer, Neuron, Target_Label)', neuron_dict) 1337 | 1338 | with open( log_dir+log_name ,'a') as f: 1339 | f.write('Compromised Neuron Candidates (Layer, Neuron, Target_Label):'+ str(neuron_dict)) 1340 | f.write('\n') 1341 | 1342 | sample_end = time.time() 1343 | 1344 | results = re_mask(model_type, model, neuron_dict, children, xs, ys, scratch_dirpath, re_epochs) 1345 | reasr_info = [] 1346 | reasrs = [] 1347 | 1348 | test_index = 0 1349 | if len(results) > 0: 1350 | reasrs = [] 1351 | test_results = [] 1352 | for result in results: 1353 | reasr,test_result = test(model, model_type, test_xs, test_ys, result, scratch_dirpath, test_index, log_dir, log_name) 1354 | reasrs.append(reasr) 1355 | test_results.append(test_result) 1356 | adv, rdelta, rmask, optz_label, RE_img, RE_mask, RE_delta, samp_label, acc = result 1357 | rmask = rmask * rmask > mask_epsilon 1358 | reasr_info.append([reasr, 'mask', str(optz_label), str(samp_label), RE_img, RE_mask, RE_delta, np.sum(rmask), acc]) 1359 | test_index = test_index + 1 1360 | print(str(model_filepath), 'mask check', max(reasrs)) 1361 | else: 1362 | print(str(model_filepath), 'mask check', 0) 1363 | 1364 | optm_end = time.time() 1365 | if len(reasrs) > 0: 1366 | freasr = max(reasrs) 1367 | f_id = reasrs.index(freasr) 1368 | else: 1369 | freasr = 0 1370 | f_id = 0 1371 | max_reasr = 0 1372 | for i in range(len(reasr_info)): 1373 | print('reasr info {0}'.format( ' '.join([str(_) for _ in reasr_info[i]]) )) 1374 | reasr = reasr_info[i][0] 1375 | if reasr > max_reasr : 1376 | max_reasr = reasr 1377 | print('{0} {1} {2} {3} {4} {5} {6} {7} {8}'.format(\ 1378 | model_filepath, model_type, 'mask', freasr, 'sampling val', sampling_val, 'time', sample_end - start, optm_end - sample_end,) ) 1379 | del model 1380 | torch.cuda.empty_cache() 1381 | return test_results -------------------------------------------------------------------------------- /data/npy_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import pickle 4 | import numpy 5 | import torchvision.transforms as transforms 6 | from PIL import Image 7 | import albumentations 8 | def addnoise(img,p=1,mean=0,var_limit=(0,1)): 9 | #img = np.asarray(img) 10 | aug = albumentations.GaussNoise(p=p,mean=mean,var_limit=var_limit) 11 | augmented = aug(image=(img).astype(numpy.uint8)) 12 | auged = augmented['image'] 13 | #auged = Image.fromarray(auged) 14 | return auged 15 | 16 | class dataset(): 17 | def __init__(self, root=None, train=True, transform=None, all_noise=False): 18 | self.root = root 19 | self.train = train 20 | self.transform = transform 21 | self.all_noise = all_noise 22 | 23 | 24 | if self.train: 25 | train_data_path = os.path.join(root, 'train_images.npy') 26 | train_labels_path = os.path.join(root, 'train_labels.npy') 27 | self.train_data = numpy.load(open(train_data_path, 'rb')) 28 | #self.train_data = torch.from_numpy(self.train_data.astype('float32')) 29 | #self.train_data = numpy.transpose(self.train_data, (0,3,1,2)) 30 | #self.train_data = Image.fromarray(self.train_data) 31 | self.train_labels = numpy.load(open(train_labels_path, 'rb')).astype('int') 32 | else: 33 | test_data_path = os.path.join(root, 'test_images.npy') 34 | test_labels_path = os.path.join(root, 'test_labels.npy') 35 | self.test_data = numpy.load(open(test_data_path, 'rb')) 36 | #self.test_data = torch.from_numpy(self.test_data.astype('float32')) 37 | #self.test_data = numpy.transpose(self.test_data, (0,3,1,2)) 38 | #self.test_data = Image.fromarray(self.test_data) 39 | self.test_labels = numpy.load(open(test_labels_path, 'rb')).astype('int') 40 | #print(self.test_data.shape) 41 | #print(self.test_labels) 42 | 43 | def __len__(self): 44 | if self.train: 45 | return len(self.train_data) 46 | else: 47 | return len(self.test_data) 48 | 49 | def __getitem__(self, index): 50 | if self.train: 51 | if self.transform: 52 | #print(self.train_data.shape) 53 | img, target = self.transform(Image.fromarray(numpy.uint8(self.train_data[index]))), self.train_labels[index] 54 | else: 55 | img, target = self.train_data[index], self.train_labels[index] 56 | else: 57 | if self.transform: 58 | #img, target = self.transform(Image.fromarray(numpy.uint8(self.test_data[index]))), self.test_labels[index] 59 | img = numpy.uint8(self.test_data[index]) 60 | if self.all_noise: 61 | img = addnoise(img,p=1,mean=80,var_limit=(0,80)) 62 | img = self.transform(Image.fromarray(img)) 63 | target = self.test_labels[index] 64 | 65 | else: 66 | img, target = self.test_data[index], self.test_labels[index] 67 | 68 | 69 | 70 | return img, target 71 | -------------------------------------------------------------------------------- /data/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import csv 4 | 5 | class HandleCsv: 6 | # 定义存放csv内容的list 7 | csv_list = [] 8 | 9 | def __init__(self, filename): 10 | self.filename = filename 11 | with open(self.filename)as fp: 12 | self.csv_list = list(csv.reader(fp)) 13 | #print(self.csv_list) 14 | 15 | # 在第N行第M列空白单元格处修改内容 16 | def modify(self, n, m, value): 17 | self.csv_list[n - 1][m - 1] = value 18 | 19 | # 插入第N行 20 | def insert_row(self, n): 21 | self.csv_list.insert(n - 1, []) 22 | 23 | # 在第N行第M列单元格插入 24 | def insert_col(self, n, m, value): 25 | # 如果该单元格左边的单元格为空，那么先对左边的单元格写入空格 26 | if len(self.csv_list[n - 1]) < m: 27 | if len(self.csv_list[n - 1]) == m - 1: 28 | self.csv_list[n - 1].append(value) 29 | else: 30 | for i in range(len(self.csv_list[n - 1]), m - 1): 31 | self.csv_list[n - 1].append('') 32 | self.csv_list[n - 1].append(value) 33 | else: 34 | self.modify(n, m, value) 35 | 36 | # 删除第N行 37 | def del_row(self, n): 38 | del self.csv_list[n - 1] 39 | 40 | # 获取第n行第m列单元格内容 41 | def get_value(self, n, m): 42 | return self.csv_list[n - 1][m - 1] 43 | 44 | def list2csv(self, file_path): 45 | try: 46 | fp = open(file_path, 'w') 47 | for items in self.csv_list: 48 | for i in range(len(items)): 49 | # 若元素中含有逗号，那么需要加双引号 50 | if items[i].find(',') != -1: 51 | fp.write('\"') 52 | fp.write(items[i]) 53 | fp.write('\"') 54 | else: 55 | fp.write(items[i]) 56 | # 最后一个元素不用加逗号 57 | if i < len(items) - 1: 58 | fp.write(',') 59 | fp.write('\n') 60 | except Exception as e: 61 | print(e) 62 | 63 | def tensor2im(input_image, imtype=np.uint8): 64 | #print(input_image.shape) 65 | mean = [0.4914, 0.4822, 0.4465] 66 | std = [0.2023, 0.1994, 0.2010] 67 | if not isinstance(input_image, np.ndarray): 68 | if isinstance(input_image, torch.Tensor): 69 | image_tensor = input_image.data 70 | else: 71 | return input_image 72 | image_numpy = image_tensor.cpu().float().numpy() 73 | if image_numpy.shape[0] == 1: 74 | image_numpy = np.tile(image_numpy, (3, 1, 1)) 75 | for i in range(len(mean)): 76 | image_numpy[i] = image_numpy[i] * std[i] + mean[i] 77 | image_numpy = image_numpy * 255 78 | image_numpy = np.transpose(image_numpy, (1, 2, 0)) 79 | else: 80 | image_numpy = input_image 81 | return image_numpy.astype(imtype) 82 | 83 | def tensor2im_gtsrb(input_image, imtype=np.uint8): 84 | #print(input_image.shape) 85 | mean = [0.3403, 0.3121, 0.3214] 86 | std = [0.2724, 0.2608, 0.2669] 87 | if not isinstance(input_image, np.ndarray): 88 | if isinstance(input_image, torch.Tensor): 89 | image_tensor = input_image.data 90 | else: 91 | return input_image 92 | image_numpy = image_tensor.cpu().float().numpy() 93 | if image_numpy.shape[0] == 1: 94 | image_numpy = np.tile(image_numpy, (3, 1, 1)) 95 | for i in range(len(mean)): 96 | image_numpy[i] = image_numpy[i] * std[i] + mean[i] 97 | image_numpy = image_numpy * 255 98 | image_numpy = np.transpose(image_numpy, (1, 2, 0)) 99 | else: 100 | image_numpy = input_image 101 | return image_numpy.astype(imtype) 102 | 103 | def tensor2im_imagenette(input_image, imtype=np.uint8): 104 | #print(input_image.shape) 105 | mean = [0.485, 0.456, 0.406] 106 | std = [0.229, 0.224, 0.225] 107 | if not isinstance(input_image, np.ndarray): 108 | if isinstance(input_image, torch.Tensor): 109 | image_tensor = input_image.data 110 | else: 111 | return input_image 112 | image_numpy = image_tensor.cpu().float().numpy() 113 | if image_numpy.shape[0] == 1: 114 | image_numpy = np.tile(image_numpy, (3, 1, 1)) 115 | for i in range(len(mean)): 116 | image_numpy[i] = image_numpy[i] * std[i] + mean[i] 117 | image_numpy = image_numpy * 255 118 | image_numpy = np.transpose(image_numpy, (1, 2, 0)) 119 | else: 120 | image_numpy = input_image 121 | return image_numpy.astype(imtype) 122 | 123 | def tensor2im_mnist(input_image, imtype=np.uint8): 124 | #print(input_image.shape) 125 | mean = [0.1307] 126 | std = [0.3081] 127 | if not isinstance(input_image, np.ndarray): 128 | if isinstance(input_image, torch.Tensor): 129 | image_tensor = input_image.data 130 | else: 131 | return input_image 132 | image_numpy = image_tensor.cpu().float().numpy() 133 | '''if image_numpy.shape[0] == 1: 134 | image_numpy = np.tile(image_numpy, (3, 1, 1))''' 135 | for i in range(len(mean)): 136 | image_numpy[i] = image_numpy[i] * std[i] + mean[i] 137 | image_numpy = image_numpy * 255 138 | image_numpy = np.transpose(image_numpy, (1, 2, 0)) 139 | else: 140 | image_numpy = input_image 141 | return image_numpy.astype(imtype) 142 | 143 | 144 | def _to_one_hot(labels,num_classes): 145 | one_hot = torch.zeros((labels.size(0), num_classes), requires_grad=True) 146 | for i, label in enumerate(labels): 147 | one_hot[i, label] = 1 148 | 149 | one_hot = torch.autograd.Variable(one_hot) 150 | 151 | return one_hot 152 | 153 | class generated_dataset(): 154 | def __init__(self, data, label): 155 | 156 | self.train_data = data 157 | self.train_labels = label 158 | 159 | def __len__(self): 160 | return len(self.train_data) 161 | 162 | def __getitem__(self, index): 163 | img, target = self.train_data[index], self.train_labels[index] 164 | 165 | return img, target -------------------------------------------------------------------------------- /generate_examples.py: -------------------------------------------------------------------------------- 1 | import torchvision.transforms as transforms 2 | import torch 3 | #print(torch.cuda.device_count()) 4 | import torchvision 5 | import torchvision.datasets as datasets 6 | import cv2 7 | import numpy as np 8 | import argparse 9 | 10 | def tensor2im(input_image, imtype=np.uint8): 11 | #print(input_image.shape) 12 | mean = [0, 0, 0] 13 | std = [1, 1, 1] 14 | if not isinstance(input_image, np.ndarray): 15 | if isinstance(input_image, torch.Tensor): 16 | image_tensor = input_image.data 17 | else: 18 | return input_image 19 | image_numpy = image_tensor.cpu().float().numpy() 20 | if image_numpy.shape[0] == 1: 21 | image_numpy = np.tile(image_numpy, (3, 1, 1)) 22 | for i in range(len(mean)): 23 | image_numpy[i] = image_numpy[i] * std[i] + mean[i] 24 | image_numpy = image_numpy * 255 25 | image_numpy = np.transpose(image_numpy, (1, 2, 0)) 26 | else: 27 | image_numpy = input_image 28 | return image_numpy.astype(imtype) 29 | 30 | parser = argparse.ArgumentParser(description='') 31 | parser.add_argument('--save_dir', default='./example/', type=str, metavar='PATH') 32 | parser.add_argument('--dataset', type=str, default="cifar10") 33 | 34 | args = parser.parse_args() 35 | 36 | 37 | if args.dataset == "cifar10": 38 | 39 | transform_test = transforms.Compose([ 40 | transforms.ToTensor(), 41 | #transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 42 | ]) 43 | 44 | dataloader = datasets.CIFAR10 45 | testset = dataloader(root='./data', train=False, download=False, transform=transform_test) 46 | testloader = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, num_workers=4) 47 | 48 | counter = [0]*10 49 | for batch_idx, (data, target) in enumerate(testloader): 50 | print(counter) 51 | cv2.imwrite(args.save_dir+"class_" + str(target.item()) + "_example_" + str(counter[int(target.item())]) + ".png", tensor2im(data[0,:,:,:])) 52 | counter[int(target.item())] = counter[int(target.item())] + 1 53 | -------------------------------------------------------------------------------- /get_data.py: -------------------------------------------------------------------------------- 1 | from data import utils,npy_data 2 | import torchvision.transforms as transforms 3 | import torchvision.datasets as datasets 4 | import torch 5 | print(torch.cuda.device_count()) 6 | def get_dataloaders(dataset_name,poison_type,trigger_size,poison_rate): 7 | 8 | if dataset_name == "cifar10": 9 | transform_train = transforms.Compose([ 10 | transforms.RandomCrop(32, padding=4), 11 | transforms.RandomHorizontalFlip(), 12 | transforms.ToTensor(), 13 | transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 14 | ]) 15 | 16 | transform_test = transforms.Compose([ 17 | transforms.ToTensor(), 18 | transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), 19 | ]) 20 | 21 | if poison_type == "none": 22 | dataloader = datasets.CIFAR10 23 | trainset = dataloader(root='./data', train=True, download=True, transform=transform_train) 24 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4) 25 | 26 | trainset_transform_test = dataloader(root='./data', train=True, download=True, transform=transform_test) 27 | trainloader_no_shuffle = torch.utils.data.DataLoader(trainset_transform_test, batch_size=128, shuffle=False, num_workers=4) 28 | trainloader_no_shuffle_bs1 = torch.utils.data.DataLoader(trainset_transform_test, batch_size=1, shuffle=False) 29 | 30 | testset = npy_data.dataset(root='./dirty_label_p05_s3', train=False, transform=transform_test) 31 | testloader_poisoned = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=4) 32 | 33 | elif poison_type == "label_specific": 34 | if poison_rate == 0.05: 35 | dataloader = datasets.CIFAR10 36 | trainset = dataloader(root='./data', train=True, download=True, transform=transform_train) 37 | trainloader_clean = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4) 38 | 39 | trainset = npy_data.dataset(root='./label_specific_p5_s3', train=True, transform=transform_train) 40 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4) 41 | 42 | trainset_transform_test = npy_data.dataset(root='./label_specific_p5_s3', train=True, transform=transform_test) 43 | trainloader_no_shuffle = torch.utils.data.DataLoader(trainset_transform_test, batch_size=128, shuffle=False, num_workers=4) 44 | trainloader_no_shuffle_bs1 = torch.utils.data.DataLoader(trainset_transform_test, batch_size=1, shuffle=False) 45 | 46 | testset = npy_data.dataset(root='./label_specific_p5_s3', train=False, transform=transform_test) 47 | testloader_poisoned = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=4) 48 | 49 | elif poison_type == "single_target": 50 | 51 | if (trigger_size == 3) and (poison_rate==0.05): 52 | trainset = npy_data.dataset(root='./dirty_label_p5_s3', train=True, transform=transform_train) 53 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4) 54 | 55 | trainset_transform_test = npy_data.dataset(root='./dirty_label_p5_s3', train=True, transform=transform_test) 56 | trainloader_no_shuffle = torch.utils.data.DataLoader(trainset_transform_test, batch_size=128, shuffle=False, num_workers=4) 57 | trainloader_no_shuffle_bs1 = torch.utils.data.DataLoader(trainset_transform_test, batch_size=1, shuffle=False) 58 | 59 | testset = npy_data.dataset(root='./dirty_label_p5_s3', train=False, transform=transform_test) 60 | testloader_poisoned = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=4) 61 | 62 | elif poison_type == "clean_label": 63 | dataloader = datasets.CIFAR10 64 | trainset = dataloader(root='./data', train=True, download=True, transform=transform_train) 65 | trainloader_clean = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4) 66 | 67 | trainset = npy_data.dataset(root='./clean_label_two1200_p5_s3', train=True, transform=transform_train)#acctually is 0.005! 68 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=4) 69 | 70 | trainset_transform_test = npy_data.dataset(root='./clean_label_two1200_p5_s3', train=True, transform=transform_test) 71 | trainloader_no_shuffle = torch.utils.data.DataLoader(trainset_transform_test, batch_size=128, shuffle=False, num_workers=4) 72 | trainloader_no_shuffle_bs1 = torch.utils.data.DataLoader(trainset_transform_test, batch_size=1, shuffle=False) 73 | 74 | testset = npy_data.dataset(root='./clean_label_two1200_p5_s3', train=False, transform=transform_test) 75 | testloader_poisoned = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=4) 76 | 77 | else: 78 | print("poison_type is not supported") 79 | 80 | 81 | dataloader = datasets.CIFAR10 82 | testset = dataloader(root='./data', train=False, download=False, transform=transform_test) 83 | testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=4) 84 | 85 | classes = ('plane', 'car', 'bird', 'cat', 86 | 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') 87 | img_size=32 88 | num_classes = 10 89 | else: 90 | print("dataset is not supported") 91 | 92 | return num_classes,trainset,trainloader,testloader,testloader_poisoned,trainloader_no_shuffle -------------------------------------------------------------------------------- /models/nin.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class Net(nn.Module): 4 | def __init__(self,num_classes=10,in_channels=3): 5 | super(Net, self).__init__() 6 | use_relu_inplace=False 7 | self.num_classes = num_classes 8 | self.classifier = nn.Sequential( 9 | nn.Conv2d(in_channels, 192, kernel_size=5, stride=1, padding=2), 10 | nn.ReLU(inplace=use_relu_inplace), 11 | nn.Conv2d(192, 160, kernel_size=1, stride=1, padding=0), 12 | nn.ReLU(inplace=use_relu_inplace), 13 | nn.Conv2d(160, 96, kernel_size=1, stride=1, padding=0), 14 | nn.ReLU(inplace=use_relu_inplace), 15 | nn.MaxPool2d(kernel_size=3, stride=2, padding=1), 16 | nn.Dropout(0.5), 17 | 18 | nn.Conv2d(96, 192, kernel_size=5, stride=1, padding=2), 19 | nn.ReLU(inplace=use_relu_inplace), 20 | nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0), 21 | nn.ReLU(inplace=use_relu_inplace), 22 | nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0), 23 | nn.ReLU(inplace=use_relu_inplace), 24 | nn.AvgPool2d(kernel_size=3, stride=2, padding=1), 25 | nn.Dropout(0.5), 26 | 27 | nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1), 28 | nn.ReLU(inplace=use_relu_inplace), 29 | nn.Conv2d(192, 192, kernel_size=1, stride=1, padding=0), 30 | nn.ReLU(inplace=use_relu_inplace), 31 | nn.Conv2d(192, num_classes, kernel_size=1, stride=1, padding=0), 32 | nn.ReLU(inplace=use_relu_inplace), 33 | nn.AvgPool2d(kernel_size=8, stride=1, padding=0), 34 | 35 | ) 36 | 37 | def forward(self, x): 38 | x = self.classifier(x) 39 | x = x.view(x.size(0), self.num_classes) 40 | return x 41 | 42 | def get_fm(self, x): 43 | for i in range(23): 44 | x = self.classifier[i](x) 45 | if i == 19: 46 | return x 47 | 48 | 49 | def get_conv_activation(self, x): 50 | for i in range(18): 51 | x = self.classifier[i](x) 52 | return x 53 | 54 | def get_all_inner_activation(self, x): 55 | inner_output_index = [0,2,4,8,10,12,16,18] 56 | inner_output_list = [] 57 | for i in range(23): 58 | x = self.classifier[i](x) 59 | if i in inner_output_index: 60 | inner_output_list.append(x) 61 | x = x.view(x.size(0), self.num_classes) 62 | return x,inner_output_list 63 | -------------------------------------------------------------------------------- /models/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | 4 | 5 | def conv3x3(in_planes, out_planes, stride=1): 6 | # 3x3 convolution with padding 7 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) 8 | 9 | 10 | class BasicBlock(nn.Module): 11 | expansion = 1 12 | 13 | def __init__(self, inplanes, planes, stride=1, downsample=None): 14 | super(BasicBlock, self).__init__() 15 | self.conv1 = conv3x3(inplanes, planes, stride) 16 | self.bn1 = nn.BatchNorm2d(planes) 17 | self.relu = nn.ReLU(inplace=True) 18 | self.conv2 = conv3x3(planes, planes) 19 | self.bn2 = nn.BatchNorm2d(planes) 20 | #print(downsample) 21 | self.downsample = downsample 22 | self.stride = stride 23 | 24 | def forward(self, x): 25 | residual = x 26 | 27 | x = self.conv1(x) 28 | x = self.bn1(x) 29 | x = self.relu(x) 30 | 31 | x = self.conv2(x) 32 | x = self.bn2(x) 33 | 34 | if self.downsample is not None: 35 | #print(x.shape) 36 | residual = self.downsample(residual) 37 | 38 | x += residual 39 | x = self.relu(x) 40 | 41 | return x 42 | 43 | def input_to_residual(self, x): 44 | residual = x 45 | if self.downsample is not None: 46 | residual = self.downsample(residual) 47 | return residual 48 | 49 | def residual_to_output(self, residual,conv2): 50 | x = residual + conv2 51 | x = self.relu(x) 52 | 53 | return x 54 | 55 | 56 | def input_to_conv2(self, x): 57 | residual = x 58 | x = self.conv1(x) 59 | x = self.bn1(x) 60 | x = self.relu(x) 61 | x = self.conv2(x) 62 | return x 63 | 64 | def conv2_to_output(self, x, residual): 65 | x = self.bn2(x) 66 | x = residual + x 67 | x = self.relu(x) 68 | return x 69 | 70 | def input_to_conv1(self, x): 71 | x = self.conv1(x) 72 | return x 73 | 74 | def conv1_to_output(self, x, residual): 75 | x = self.bn1(x) 76 | x = self.relu(x) 77 | 78 | x = self.conv2(x) 79 | x = self.bn2(x) 80 | 81 | x += residual 82 | x = self.relu(x) 83 | 84 | return x 85 | 86 | class Bottleneck(nn.Module): 87 | expansion = 4 88 | 89 | def __init__(self, inplanes, planes, stride=1, downsample=None): 90 | super(Bottleneck, self).__init__() 91 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 92 | self.bn1 = nn.BatchNorm2d(planes) 93 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 94 | self.bn2 = nn.BatchNorm2d(planes) 95 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 96 | self.bn3 = nn.BatchNorm2d(planes * 4) 97 | self.relu = nn.ReLU(inplace=True) 98 | self.downsample = downsample 99 | self.stride = stride 100 | 101 | def forward(self, x): 102 | residual = x 103 | 104 | x = self.conv1(x) 105 | x = self.bn1(x) 106 | x = self.relu(x) 107 | 108 | x = self.conv2(x) 109 | x = self.bn2(x) 110 | x = self.relu(x) 111 | 112 | x = self.conv3(x) 113 | x = self.bn3(x) 114 | 115 | if self.downsample is not None: 116 | residual = self.downsample(residual) 117 | 118 | x += residual 119 | x = self.relu(x) 120 | 121 | return x 122 | 123 | 124 | class ResNet(nn.Module): 125 | 126 | def __init__(self, block, layers, num_classes=10,in_channels=3): 127 | self.inplanes = 64 128 | super(ResNet, self).__init__() 129 | self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1, bias=False) 130 | self.bn1 = nn.BatchNorm2d(64) 131 | self.relu = nn.ReLU(inplace=True) 132 | self.layer1 = self._make_layer(block, 64, layers[0]) 133 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 134 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 135 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 136 | self.avgpool = nn.AvgPool2d(kernel_size=4) 137 | self.fc = nn.Linear(512 * block.expansion, num_classes) 138 | 139 | self.inter_feature = {} 140 | self.inter_gradient = {} 141 | 142 | self.register_all_hooks() 143 | 144 | for m in self.modules(): 145 | if isinstance(m, nn.Conv2d): 146 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 147 | m.weight.data.normal_(0, math.sqrt(2. / n)) 148 | elif isinstance(m, nn.BatchNorm2d): 149 | m.weight.data.fill_(1) 150 | m.bias.data.zero_() 151 | 152 | def _make_layer(self, block, planes, blocks, stride=1): 153 | downsample = None 154 | if stride != 1 or self.inplanes != planes * block.expansion: 155 | downsample = nn.Sequential( 156 | nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), 157 | nn.BatchNorm2d(planes * block.expansion), 158 | ) 159 | 160 | layers = [] 161 | layers.append(block(self.inplanes, planes, stride, downsample)) 162 | self.inplanes = planes * block.expansion 163 | for i in range(1, blocks): 164 | layers.append(block(self.inplanes, planes)) 165 | return nn.Sequential(*layers) 166 | 167 | def forward(self, x): 168 | x = self.conv1(x) 169 | x = self.bn1(x) 170 | x = self.relu(x) 171 | 172 | x = self.layer1(x) 173 | x = self.layer2(x) 174 | x = self.layer3(x) 175 | x = self.layer4(x) 176 | 177 | x = self.avgpool(x) 178 | x = x.view(x.size(0), -1) 179 | x = self.fc(x) 180 | 181 | return x 182 | 183 | def get_fm(self, x): 184 | x = self.conv1(x) 185 | x = self.bn1(x) 186 | x = self.relu(x) 187 | 188 | x = self.layer1(x) 189 | x = self.layer2(x) 190 | x = self.layer3(x) 191 | x = self.layer4(x) 192 | 193 | #x = self.avgpool(x) 194 | 195 | return x 196 | 197 | def make_hook(self, name, flag): 198 | if flag == 'forward': 199 | def hook(m, input, output): 200 | self.inter_feature[name] = output 201 | return hook 202 | elif flag == 'backward': 203 | def hook(m, input, output): 204 | self.inter_gradient[name] = output 205 | return hook 206 | else: 207 | assert False 208 | 209 | def register_all_hooks(self): 210 | self.conv1.register_forward_hook(self.make_hook("Conv1_Conv1_Conv1_", 'forward')) 211 | self.layer1[0].conv1.register_forward_hook(self.make_hook("Layer1_0_Conv1_", 'forward')) 212 | self.layer1[0].conv2.register_forward_hook(self.make_hook("Layer1_0_Conv2_", 'forward')) 213 | self.layer1[1].conv1.register_forward_hook(self.make_hook("Layer1_1_Conv1_", 'forward')) 214 | self.layer1[1].conv2.register_forward_hook(self.make_hook("Layer1_1_Conv2_", 'forward')) 215 | 216 | self.layer2[0].conv1.register_forward_hook(self.make_hook("Layer2_0_Conv1_", 'forward')) 217 | self.layer2[0].downsample.register_forward_hook(self.make_hook("Layer2_0_Downsample_", 'forward')) 218 | self.layer2[0].conv2.register_forward_hook(self.make_hook("Layer2_0_Conv2_", 'forward')) 219 | self.layer2[1].conv1.register_forward_hook(self.make_hook("Layer2_1_Conv1_", 'forward')) 220 | self.layer2[1].conv2.register_forward_hook(self.make_hook("Layer2_1_Conv2_", 'forward')) 221 | 222 | self.layer3[0].conv1.register_forward_hook(self.make_hook("Layer3_0_Conv1_", 'forward')) 223 | self.layer3[0].downsample.register_forward_hook(self.make_hook("Layer3_0_Downsample_", 'forward')) 224 | self.layer3[0].conv2.register_forward_hook(self.make_hook("Layer3_0_Conv2_", 'forward')) 225 | self.layer3[1].conv1.register_forward_hook(self.make_hook("Layer3_1_Conv1_", 'forward')) 226 | self.layer3[1].conv2.register_forward_hook(self.make_hook("Layer3_1_Conv2_", 'forward')) 227 | 228 | self.layer4[0].conv1.register_forward_hook(self.make_hook("Layer4_0_Conv1_", 'forward')) 229 | self.layer4[0].downsample.register_forward_hook(self.make_hook("Layer4_0_Downsample_", 'forward')) 230 | self.layer4[0].conv2.register_forward_hook(self.make_hook("Layer4_0_Conv2_", 'forward')) 231 | self.layer4[1].conv1.register_forward_hook(self.make_hook("Layer4_1_Conv1_", 'forward')) 232 | self.layer4[1].conv2.register_forward_hook(self.make_hook("Layer4_1_Conv2_", 'forward')) 233 | 234 | 235 | 236 | '''def get_all_inner_activation(self, x): 237 | inner_output_index = [0,2,4,8,10,12,16,18] 238 | inner_output_list = [] 239 | for i in range(23): 240 | x = self.classifier[i](x) 241 | if i in inner_output_index: 242 | inner_output_list.append(x) 243 | x = x.view(x.size(0), self.num_classes) 244 | return x,inner_output_list''' 245 | 246 | ############################################################################# 247 | def input_to_conv1(self, x): 248 | x = self.conv1(x) 249 | return x 250 | 251 | def conv1_to_output(self, x): 252 | x = self.bn1(x) 253 | x = self.relu(x) 254 | 255 | x = self.layer1(x) 256 | x = self.layer2(x) 257 | x = self.layer3(x) 258 | x = self.layer4(x) 259 | 260 | x = self.avgpool(x) 261 | x = x.view(x.size(0), -1) 262 | x = self.fc(x) 263 | 264 | return x 265 | 266 | ############################################################################# 267 | def input_to_layer1_0_residual(self, x): 268 | x = self.conv1(x) 269 | x = self.bn1(x) 270 | x = self.relu(x) 271 | 272 | x = self.layer1[0].input_to_residual(x) 273 | 274 | return x 275 | 276 | def layer1_0_residual_to_output(self, residual, conv2): 277 | 278 | x = self.layer1[0].residual_to_output(residual,conv2) 279 | x = self.layer1[1](x) 280 | x = self.layer2(x) 281 | x = self.layer3(x) 282 | x = self.layer4(x) 283 | 284 | x = self.avgpool(x) 285 | x = x.view(x.size(0), -1) 286 | x = self.fc(x) 287 | return x 288 | 289 | def input_to_layer1_0_conv2(self, x): 290 | x = self.conv1(x) 291 | x = self.bn1(x) 292 | x = self.relu(x) 293 | x = self.layer1[0].input_to_conv2(x) 294 | return x 295 | 296 | def layer1_0_conv2_to_output(self, x, residual): 297 | x = self.layer1[0].conv2_to_output(x, residual) 298 | x = self.layer1[1](x) 299 | x = self.layer2(x) 300 | x = self.layer3(x) 301 | x = self.layer4(x) 302 | x = self.avgpool(x) 303 | x = x.view(x.size(0), -1) 304 | x = self.fc(x) 305 | return x 306 | 307 | def input_to_layer1_0_conv1(self, x): 308 | x = self.conv1(x) 309 | x = self.bn1(x) 310 | x = self.relu(x) 311 | x = self.layer1[0].input_to_conv1(x) 312 | return x 313 | 314 | def layer1_0_conv1_to_output(self, x, residual): 315 | x = self.layer1[0].conv1_to_output(x, residual) 316 | x = self.layer1[1](x) 317 | x = self.layer2(x) 318 | x = self.layer3(x) 319 | x = self.layer4(x) 320 | x = self.avgpool(x) 321 | x = x.view(x.size(0), -1) 322 | x = self.fc(x) 323 | return x 324 | ############################################################################# 325 | 326 | def input_to_layer1_1_residual(self, x): 327 | x = self.conv1(x) 328 | x = self.bn1(x) 329 | x = self.relu(x) 330 | x = self.layer1[0](x) 331 | x = self.layer1[1].input_to_residual(x) 332 | 333 | return x 334 | 335 | def input_to_layer1_1_conv2(self, x): 336 | x = self.conv1(x) 337 | x = self.bn1(x) 338 | x = self.relu(x) 339 | x = self.layer1[0](x) 340 | x = self.layer1[1].input_to_conv2(x) 341 | return x 342 | 343 | def layer1_1_conv2_to_output(self, x, residual): 344 | x = self.layer1[1].conv2_to_output(x, residual) 345 | x = self.layer2(x) 346 | x = self.layer3(x) 347 | x = self.layer4(x) 348 | x = self.avgpool(x) 349 | x = x.view(x.size(0), -1) 350 | x = self.fc(x) 351 | return x 352 | 353 | def input_to_layer1_1_conv1(self, x): 354 | x = self.conv1(x) 355 | x = self.bn1(x) 356 | x = self.relu(x) 357 | x = self.layer1[0](x) 358 | x = self.layer1[1].input_to_conv1(x) 359 | return x 360 | 361 | def layer1_1_conv1_to_output(self, x, residual): 362 | x = self.layer1[1].conv1_to_output(x, residual) 363 | x = self.layer2(x) 364 | x = self.layer3(x) 365 | x = self.layer4(x) 366 | x = self.avgpool(x) 367 | x = x.view(x.size(0), -1) 368 | x = self.fc(x) 369 | return x 370 | 371 | ############################################################################# 372 | 373 | ############################################################################# 374 | def input_to_layer2_0_residual(self, x): 375 | x = self.conv1(x) 376 | x = self.bn1(x) 377 | x = self.relu(x) 378 | 379 | x = self.layer1(x) 380 | x = self.layer2[0].input_to_residual(x) 381 | 382 | return x 383 | 384 | def layer2_0_residual_to_output(self, residual, conv2): 385 | 386 | x = self.layer2[0].residual_to_output(residual,conv2) 387 | x = self.layer2[1](x) 388 | x = self.layer3(x) 389 | x = self.layer4(x) 390 | 391 | x = self.avgpool(x) 392 | x = x.view(x.size(0), -1) 393 | x = self.fc(x) 394 | return x 395 | 396 | def input_to_layer2_0_conv2(self, x): 397 | x = self.conv1(x) 398 | x = self.bn1(x) 399 | x = self.relu(x) 400 | x = self.layer1(x) 401 | x = self.layer2[0].input_to_conv2(x) 402 | return x 403 | 404 | def layer2_0_conv2_to_output(self, x, residual): 405 | x = self.layer2[0].conv2_to_output(x, residual) 406 | x = self.layer2[1](x) 407 | x = self.layer3(x) 408 | x = self.layer4(x) 409 | x = self.avgpool(x) 410 | x = x.view(x.size(0), -1) 411 | x = self.fc(x) 412 | return x 413 | 414 | def input_to_layer2_0_conv1(self, x): 415 | x = self.conv1(x) 416 | x = self.bn1(x) 417 | x = self.relu(x) 418 | x = self.layer1(x) 419 | x = self.layer2[0].input_to_conv1(x) 420 | return x 421 | 422 | def layer2_0_conv1_to_output(self, x, residual): 423 | x = self.layer2[0].conv1_to_output(x, residual) 424 | x = self.layer2[1](x) 425 | x = self.layer3(x) 426 | x = self.layer4(x) 427 | x = self.avgpool(x) 428 | x = x.view(x.size(0), -1) 429 | x = self.fc(x) 430 | return x 431 | ############################################################################# 432 | 433 | def input_to_layer2_1_residual(self, x): 434 | x = self.conv1(x) 435 | x = self.bn1(x) 436 | x = self.relu(x) 437 | 438 | x = self.layer1(x) 439 | x = self.layer2[0](x) 440 | x = self.layer2[1].input_to_residual(x) 441 | 442 | return x 443 | 444 | def input_to_layer2_1_conv2(self, x): 445 | x = self.conv1(x) 446 | x = self.bn1(x) 447 | x = self.relu(x) 448 | x = self.layer1(x) 449 | x = self.layer2[0](x) 450 | x = self.layer2[1].input_to_conv2(x) 451 | return x 452 | 453 | def layer2_1_conv2_to_output(self, x, residual): 454 | x = self.layer2[1].conv2_to_output(x, residual) 455 | x = self.layer3(x) 456 | x = self.layer4(x) 457 | x = self.avgpool(x) 458 | x = x.view(x.size(0), -1) 459 | x = self.fc(x) 460 | return x 461 | 462 | def input_to_layer2_1_conv1(self, x): 463 | x = self.conv1(x) 464 | x = self.bn1(x) 465 | x = self.relu(x) 466 | x = self.layer1(x) 467 | x = self.layer2[0](x) 468 | x = self.layer2[1].input_to_conv1(x) 469 | return x 470 | 471 | def layer2_1_conv1_to_output(self, x, residual): 472 | x = self.layer2[1].conv1_to_output(x, residual) 473 | x = self.layer3(x) 474 | x = self.layer4(x) 475 | x = self.avgpool(x) 476 | x = x.view(x.size(0), -1) 477 | x = self.fc(x) 478 | return x 479 | 480 | ############################################################################# 481 | 482 | ############################################################################# 483 | def input_to_layer3_0_residual(self, x): 484 | x = self.conv1(x) 485 | x = self.bn1(x) 486 | x = self.relu(x) 487 | 488 | x = self.layer1(x) 489 | x = self.layer2(x) 490 | x = self.layer3[0].input_to_residual(x) 491 | 492 | return x 493 | 494 | def layer3_0_residual_to_output(self, residual, conv2): 495 | 496 | x = self.layer3[0].residual_to_output(residual,conv2) 497 | x = self.layer3[1](x) 498 | x = self.layer4(x) 499 | 500 | x = self.avgpool(x) 501 | x = x.view(x.size(0), -1) 502 | x = self.fc(x) 503 | return x 504 | 505 | def input_to_layer3_0_conv2(self, x): 506 | x = self.conv1(x) 507 | x = self.bn1(x) 508 | x = self.relu(x) 509 | x = self.layer1(x) 510 | x = self.layer2(x) 511 | x = self.layer3[0].input_to_conv2(x) 512 | return x 513 | 514 | def layer3_0_conv2_to_output(self, x, residual): 515 | x = self.layer3[0].conv2_to_output(x, residual) 516 | x = self.layer3[1](x) 517 | x = self.layer4(x) 518 | x = self.avgpool(x) 519 | x = x.view(x.size(0), -1) 520 | x = self.fc(x) 521 | return x 522 | 523 | def input_to_layer3_0_conv1(self, x): 524 | x = self.conv1(x) 525 | x = self.bn1(x) 526 | x = self.relu(x) 527 | x = self.layer1(x) 528 | x = self.layer2(x) 529 | x = self.layer3[0].input_to_conv1(x) 530 | return x 531 | 532 | def layer3_0_conv1_to_output(self, x, residual): 533 | x = self.layer3[0].conv1_to_output(x, residual) 534 | x = self.layer3[1](x) 535 | x = self.layer4(x) 536 | x = self.avgpool(x) 537 | x = x.view(x.size(0), -1) 538 | x = self.fc(x) 539 | return x 540 | ############################################################################# 541 | 542 | def input_to_layer3_1_residual(self, x): 543 | x = self.conv1(x) 544 | x = self.bn1(x) 545 | x = self.relu(x) 546 | 547 | x = self.layer1(x) 548 | x = self.layer2(x) 549 | x = self.layer3[0](x) 550 | x = self.layer3[1].input_to_residual(x) 551 | 552 | return x 553 | 554 | def input_to_layer3_1_conv2(self, x): 555 | x = self.conv1(x) 556 | x = self.bn1(x) 557 | x = self.relu(x) 558 | x = self.layer1(x) 559 | x = self.layer2(x) 560 | x = self.layer3[0](x) 561 | x = self.layer3[1].input_to_conv2(x) 562 | return x 563 | 564 | def layer3_1_conv2_to_output(self, x, residual): 565 | x = self.layer3[1].conv2_to_output(x, residual) 566 | x = self.layer4(x) 567 | x = self.avgpool(x) 568 | x = x.view(x.size(0), -1) 569 | x = self.fc(x) 570 | return x 571 | 572 | def input_to_layer3_1_conv1(self, x): 573 | x = self.conv1(x) 574 | x = self.bn1(x) 575 | x = self.relu(x) 576 | x = self.layer1(x) 577 | x = self.layer2(x) 578 | x = self.layer3[0](x) 579 | x = self.layer3[1].input_to_conv1(x) 580 | return x 581 | 582 | def layer3_1_conv1_to_output(self, x, residual): 583 | x = self.layer3[1].conv1_to_output(x, residual) 584 | x = self.layer4(x) 585 | x = self.avgpool(x) 586 | x = x.view(x.size(0), -1) 587 | x = self.fc(x) 588 | return x 589 | 590 | ############################################################################# 591 | def input_to_layer4_0_residual(self, x): 592 | x = self.conv1(x) 593 | x = self.bn1(x) 594 | x = self.relu(x) 595 | 596 | x = self.layer1(x) 597 | x = self.layer2(x) 598 | x = self.layer3(x) 599 | x = self.layer4[0].input_to_residual(x) 600 | 601 | return x 602 | 603 | def layer4_0_residual_to_output(self, residual, conv2): 604 | 605 | x = self.layer4[0].residual_to_output(residual,conv2) 606 | x = self.layer4[1](x) 607 | 608 | x = self.avgpool(x) 609 | x = x.view(x.size(0), -1) 610 | x = self.fc(x) 611 | return x 612 | 613 | def input_to_layer4_0_conv2(self, x): 614 | x = self.conv1(x) 615 | x = self.bn1(x) 616 | x = self.relu(x) 617 | x = self.layer1(x) 618 | x = self.layer2(x) 619 | x = self.layer3(x) 620 | x = self.layer4[0].input_to_conv2(x) 621 | return x 622 | 623 | def layer4_0_conv2_to_output(self, x, residual): 624 | x = self.layer4[0].conv2_to_output(x, residual) 625 | x = self.layer4[1](x) 626 | x = self.avgpool(x) 627 | x = x.view(x.size(0), -1) 628 | x = self.fc(x) 629 | return x 630 | 631 | def input_to_layer4_0_conv1(self, x): 632 | x = self.conv1(x) 633 | x = self.bn1(x) 634 | x = self.relu(x) 635 | x = self.layer1(x) 636 | x = self.layer2(x) 637 | x = self.layer3(x) 638 | x = self.layer4[0].input_to_conv1(x) 639 | return x 640 | 641 | def layer4_0_conv1_to_output(self, x, residual): 642 | x = self.layer4[0].conv1_to_output(x, residual) 643 | x = self.layer4[1](x) 644 | x = self.avgpool(x) 645 | x = x.view(x.size(0), -1) 646 | x = self.fc(x) 647 | return x 648 | ############################################################################# 649 | def input_to_layer4_1_residual(self, x): 650 | x = self.conv1(x) 651 | x = self.bn1(x) 652 | x = self.relu(x) 653 | 654 | x = self.layer1(x) 655 | x = self.layer2(x) 656 | x = self.layer3(x) 657 | x = self.layer4[0](x) 658 | x = self.layer4[1].input_to_residual(x) 659 | 660 | return x 661 | 662 | def input_to_layer4_1_conv2(self, x): 663 | x = self.conv1(x) 664 | x = self.bn1(x) 665 | x = self.relu(x) 666 | x = self.layer1(x) 667 | x = self.layer2(x) 668 | x = self.layer3(x) 669 | x = self.layer4[0](x) 670 | x = self.layer4[1].input_to_conv2(x) 671 | return x 672 | 673 | def layer4_1_conv2_to_output(self, x, residual): 674 | x = self.layer4[1].conv2_to_output(x, residual) 675 | x = self.avgpool(x) 676 | x = x.view(x.size(0), -1) 677 | x = self.fc(x) 678 | return x 679 | 680 | def input_to_layer4_1_conv1(self, x): 681 | x = self.conv1(x) 682 | x = self.bn1(x) 683 | x = self.relu(x) 684 | x = self.layer1(x) 685 | x = self.layer2(x) 686 | x = self.layer3(x) 687 | x = self.layer4[0](x) 688 | x = self.layer4[1].input_to_conv1(x) 689 | return x 690 | 691 | def layer4_1_conv1_to_output(self, x, residual): 692 | x = self.layer4[1].conv1_to_output(x, residual) 693 | x = self.avgpool(x) 694 | x = x.view(x.size(0), -1) 695 | x = self.fc(x) 696 | return x 697 | ############################################################################# 698 | 699 | def resnet18(**kwargs): 700 | return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 701 | 702 | 703 | def resnet34(**kwargs): 704 | return ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 705 | 706 | 707 | def resnet50(**kwargs): 708 | return ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 709 | 710 | 711 | def resnet101(**kwargs): 712 | return ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 713 | 714 | 715 | def resnet152(**kwargs): 716 | return ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) -------------------------------------------------------------------------------- /models/resnet_more_activation.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | 4 | 5 | def conv3x3(in_planes, out_planes, stride=1): 6 | # 3x3 convolution with padding 7 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) 8 | 9 | activation_function_global = None 10 | 11 | class BasicBlock(nn.Module): 12 | expansion = 1 13 | 14 | def __init__(self, inplanes, planes, stride=1, downsample=None): 15 | super(BasicBlock, self).__init__() 16 | self.conv1 = conv3x3(inplanes, planes, stride) 17 | self.bn1 = nn.BatchNorm2d(planes) 18 | self.relu = activation_function_global 19 | self.conv2 = conv3x3(planes, planes) 20 | self.bn2 = nn.BatchNorm2d(planes) 21 | #print(downsample) 22 | self.downsample = downsample 23 | self.stride = stride 24 | 25 | def forward(self, x): 26 | residual = x 27 | 28 | x = self.conv1(x) 29 | x = self.bn1(x) 30 | x = self.relu(x) 31 | 32 | x = self.conv2(x) 33 | x = self.bn2(x) 34 | 35 | if self.downsample is not None: 36 | #print(x.shape) 37 | residual = self.downsample(residual) 38 | 39 | x += residual 40 | x = self.relu(x) 41 | 42 | return x 43 | 44 | def input_to_residual(self, x): 45 | residual = x 46 | if self.downsample is not None: 47 | residual = self.downsample(residual) 48 | return residual 49 | 50 | def residual_to_output(self, residual,conv2): 51 | x = residual + conv2 52 | x = self.relu(x) 53 | 54 | return x 55 | 56 | 57 | def input_to_conv2(self, x): 58 | residual = x 59 | x = self.conv1(x) 60 | x = self.bn1(x) 61 | x = self.relu(x) 62 | x = self.conv2(x) 63 | return x 64 | 65 | def conv2_to_output(self, x, residual): 66 | x = self.bn2(x) 67 | x = residual + x 68 | x = self.relu(x) 69 | return x 70 | 71 | def input_to_conv1(self, x): 72 | x = self.conv1(x) 73 | return x 74 | 75 | def conv1_to_output(self, x, residual): 76 | x = self.bn1(x) 77 | x = self.relu(x) 78 | 79 | x = self.conv2(x) 80 | x = self.bn2(x) 81 | 82 | x += residual 83 | x = self.relu(x) 84 | 85 | return x 86 | 87 | class Bottleneck(nn.Module): 88 | expansion = 4 89 | 90 | def __init__(self, inplanes, planes, stride=1, downsample=None): 91 | super(Bottleneck, self).__init__() 92 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 93 | self.bn1 = nn.BatchNorm2d(planes) 94 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 95 | self.bn2 = nn.BatchNorm2d(planes) 96 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 97 | self.bn3 = nn.BatchNorm2d(planes * 4) 98 | self.relu = activation_function_global 99 | self.downsample = downsample 100 | self.stride = stride 101 | 102 | def forward(self, x): 103 | residual = x 104 | 105 | x = self.conv1(x) 106 | x = self.bn1(x) 107 | x = self.relu(x) 108 | 109 | x = self.conv2(x) 110 | x = self.bn2(x) 111 | x = self.relu(x) 112 | 113 | x = self.conv3(x) 114 | x = self.bn3(x) 115 | 116 | if self.downsample is not None: 117 | residual = self.downsample(residual) 118 | 119 | x += residual 120 | x = self.relu(x) 121 | 122 | return x 123 | 124 | 125 | class ResNet(nn.Module): 126 | 127 | def __init__(self, block, layers, num_classes=10,in_channels=3, activation_function=None): 128 | 129 | global activation_function_global 130 | 131 | if activation_function=="leakyrelu": 132 | activation_function_global = nn.LeakyReLU(inplace=True) 133 | elif activation_function=="elu": 134 | activation_function_global = nn.ELU(inplace=True) 135 | elif activation_function=="prelu": 136 | activation_function_global = nn.PReLU() 137 | elif activation_function=="tanhshrink": 138 | activation_function_global = nn.Tanhshrink() 139 | elif activation_function=="softplus": 140 | activation_function_global = nn.Softplus() 141 | 142 | self.inplanes = 64 143 | super(ResNet, self).__init__() 144 | self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1, bias=False) 145 | self.bn1 = nn.BatchNorm2d(64) 146 | self.relu = activation_function_global 147 | self.layer1 = self._make_layer(block, 64, layers[0]) 148 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 149 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 150 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 151 | self.avgpool = nn.AvgPool2d(kernel_size=4) 152 | self.fc = nn.Linear(512 * block.expansion, num_classes) 153 | 154 | self.inter_feature = {} 155 | self.inter_gradient = {} 156 | 157 | self.register_all_hooks() 158 | 159 | for m in self.modules(): 160 | if isinstance(m, nn.Conv2d): 161 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 162 | m.weight.data.normal_(0, math.sqrt(2. / n)) 163 | elif isinstance(m, nn.BatchNorm2d): 164 | m.weight.data.fill_(1) 165 | m.bias.data.zero_() 166 | 167 | def _make_layer(self, block, planes, blocks, stride=1): 168 | downsample = None 169 | if stride != 1 or self.inplanes != planes * block.expansion: 170 | downsample = nn.Sequential( 171 | nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), 172 | nn.BatchNorm2d(planes * block.expansion), 173 | ) 174 | 175 | layers = [] 176 | layers.append(block(self.inplanes, planes, stride, downsample)) 177 | self.inplanes = planes * block.expansion 178 | for i in range(1, blocks): 179 | layers.append(block(self.inplanes, planes)) 180 | return nn.Sequential(*layers) 181 | 182 | def forward(self, x): 183 | x = self.conv1(x) 184 | x = self.bn1(x) 185 | x = self.relu(x) 186 | 187 | x = self.layer1(x) 188 | x = self.layer2(x) 189 | x = self.layer3(x) 190 | x = self.layer4(x) 191 | 192 | x = self.avgpool(x) 193 | x = x.view(x.size(0), -1) 194 | x = self.fc(x) 195 | 196 | return x 197 | 198 | def make_hook(self, name, flag): 199 | if flag == 'forward': 200 | def hook(m, input, output): 201 | self.inter_feature[name] = output 202 | return hook 203 | elif flag == 'backward': 204 | def hook(m, input, output): 205 | self.inter_gradient[name] = output 206 | return hook 207 | else: 208 | assert False 209 | 210 | def register_all_hooks(self): 211 | self.conv1.register_forward_hook(self.make_hook("Conv1_Conv1_Conv1_", 'forward')) 212 | self.layer1[0].conv1.register_forward_hook(self.make_hook("Layer1_0_Conv1_", 'forward')) 213 | self.layer1[0].conv2.register_forward_hook(self.make_hook("Layer1_0_Conv2_", 'forward')) 214 | self.layer1[1].conv1.register_forward_hook(self.make_hook("Layer1_1_Conv1_", 'forward')) 215 | self.layer1[1].conv2.register_forward_hook(self.make_hook("Layer1_1_Conv2_", 'forward')) 216 | 217 | self.layer2[0].conv1.register_forward_hook(self.make_hook("Layer2_0_Conv1_", 'forward')) 218 | self.layer2[0].downsample.register_forward_hook(self.make_hook("Layer2_0_Downsample_", 'forward')) 219 | self.layer2[0].conv2.register_forward_hook(self.make_hook("Layer2_0_Conv2_", 'forward')) 220 | self.layer2[1].conv1.register_forward_hook(self.make_hook("Layer2_1_Conv1_", 'forward')) 221 | self.layer2[1].conv2.register_forward_hook(self.make_hook("Layer2_1_Conv2_", 'forward')) 222 | 223 | self.layer3[0].conv1.register_forward_hook(self.make_hook("Layer3_0_Conv1_", 'forward')) 224 | self.layer3[0].downsample.register_forward_hook(self.make_hook("Layer3_0_Downsample_", 'forward')) 225 | self.layer3[0].conv2.register_forward_hook(self.make_hook("Layer3_0_Conv2_", 'forward')) 226 | self.layer3[1].conv1.register_forward_hook(self.make_hook("Layer3_1_Conv1_", 'forward')) 227 | self.layer3[1].conv2.register_forward_hook(self.make_hook("Layer3_1_Conv2_", 'forward')) 228 | 229 | self.layer4[0].conv1.register_forward_hook(self.make_hook("Layer4_0_Conv1_", 'forward')) 230 | self.layer4[0].downsample.register_forward_hook(self.make_hook("Layer4_0_Downsample_", 'forward')) 231 | self.layer4[0].conv2.register_forward_hook(self.make_hook("Layer4_0_Conv2_", 'forward')) 232 | self.layer4[1].conv1.register_forward_hook(self.make_hook("Layer4_1_Conv1_", 'forward')) 233 | self.layer4[1].conv2.register_forward_hook(self.make_hook("Layer4_1_Conv2_", 'forward')) 234 | 235 | 236 | 237 | '''def get_all_inner_activation(self, x): 238 | inner_output_index = [0,2,4,8,10,12,16,18] 239 | inner_output_list = [] 240 | for i in range(23): 241 | x = self.classifier[i](x) 242 | if i in inner_output_index: 243 | inner_output_list.append(x) 244 | x = x.view(x.size(0), self.num_classes) 245 | return x,inner_output_list''' 246 | 247 | ############################################################################# 248 | def input_to_conv1(self, x): 249 | x = self.conv1(x) 250 | return x 251 | 252 | def conv1_to_output(self, x): 253 | x = self.bn1(x) 254 | x = self.relu(x) 255 | 256 | x = self.layer1(x) 257 | x = self.layer2(x) 258 | x = self.layer3(x) 259 | x = self.layer4(x) 260 | 261 | x = self.avgpool(x) 262 | x = x.view(x.size(0), -1) 263 | x = self.fc(x) 264 | 265 | return x 266 | 267 | ############################################################################# 268 | def input_to_layer1_0_residual(self, x): 269 | x = self.conv1(x) 270 | x = self.bn1(x) 271 | x = self.relu(x) 272 | 273 | x = self.layer1[0].input_to_residual(x) 274 | 275 | return x 276 | 277 | def layer1_0_residual_to_output(self, residual, conv2): 278 | 279 | x = self.layer1[0].residual_to_output(residual,conv2) 280 | x = self.layer1[1](x) 281 | x = self.layer2(x) 282 | x = self.layer3(x) 283 | x = self.layer4(x) 284 | 285 | x = self.avgpool(x) 286 | x = x.view(x.size(0), -1) 287 | x = self.fc(x) 288 | return x 289 | 290 | def input_to_layer1_0_conv2(self, x): 291 | x = self.conv1(x) 292 | x = self.bn1(x) 293 | x = self.relu(x) 294 | x = self.layer1[0].input_to_conv2(x) 295 | return x 296 | 297 | def layer1_0_conv2_to_output(self, x, residual): 298 | x = self.layer1[0].conv2_to_output(x, residual) 299 | x = self.layer1[1](x) 300 | x = self.layer2(x) 301 | x = self.layer3(x) 302 | x = self.layer4(x) 303 | x = self.avgpool(x) 304 | x = x.view(x.size(0), -1) 305 | x = self.fc(x) 306 | return x 307 | 308 | def input_to_layer1_0_conv1(self, x): 309 | x = self.conv1(x) 310 | x = self.bn1(x) 311 | x = self.relu(x) 312 | x = self.layer1[0].input_to_conv1(x) 313 | return x 314 | 315 | def layer1_0_conv1_to_output(self, x, residual): 316 | x = self.layer1[0].conv1_to_output(x, residual) 317 | x = self.layer1[1](x) 318 | x = self.layer2(x) 319 | x = self.layer3(x) 320 | x = self.layer4(x) 321 | x = self.avgpool(x) 322 | x = x.view(x.size(0), -1) 323 | x = self.fc(x) 324 | return x 325 | ############################################################################# 326 | 327 | def input_to_layer1_1_residual(self, x): 328 | x = self.conv1(x) 329 | x = self.bn1(x) 330 | x = self.relu(x) 331 | x = self.layer1[0](x) 332 | x = self.layer1[1].input_to_residual(x) 333 | 334 | return x 335 | 336 | def input_to_layer1_1_conv2(self, x): 337 | x = self.conv1(x) 338 | x = self.bn1(x) 339 | x = self.relu(x) 340 | x = self.layer1[0](x) 341 | x = self.layer1[1].input_to_conv2(x) 342 | return x 343 | 344 | def layer1_1_conv2_to_output(self, x, residual): 345 | x = self.layer1[1].conv2_to_output(x, residual) 346 | x = self.layer2(x) 347 | x = self.layer3(x) 348 | x = self.layer4(x) 349 | x = self.avgpool(x) 350 | x = x.view(x.size(0), -1) 351 | x = self.fc(x) 352 | return x 353 | 354 | def input_to_layer1_1_conv1(self, x): 355 | x = self.conv1(x) 356 | x = self.bn1(x) 357 | x = self.relu(x) 358 | x = self.layer1[0](x) 359 | x = self.layer1[1].input_to_conv1(x) 360 | return x 361 | 362 | def layer1_1_conv1_to_output(self, x, residual): 363 | x = self.layer1[1].conv1_to_output(x, residual) 364 | x = self.layer2(x) 365 | x = self.layer3(x) 366 | x = self.layer4(x) 367 | x = self.avgpool(x) 368 | x = x.view(x.size(0), -1) 369 | x = self.fc(x) 370 | return x 371 | 372 | ############################################################################# 373 | 374 | ############################################################################# 375 | def input_to_layer2_0_residual(self, x): 376 | x = self.conv1(x) 377 | x = self.bn1(x) 378 | x = self.relu(x) 379 | 380 | x = self.layer1(x) 381 | x = self.layer2[0].input_to_residual(x) 382 | 383 | return x 384 | 385 | def layer2_0_residual_to_output(self, residual, conv2): 386 | 387 | x = self.layer2[0].residual_to_output(residual,conv2) 388 | x = self.layer2[1](x) 389 | x = self.layer3(x) 390 | x = self.layer4(x) 391 | 392 | x = self.avgpool(x) 393 | x = x.view(x.size(0), -1) 394 | x = self.fc(x) 395 | return x 396 | 397 | def input_to_layer2_0_conv2(self, x): 398 | x = self.conv1(x) 399 | x = self.bn1(x) 400 | x = self.relu(x) 401 | x = self.layer1(x) 402 | x = self.layer2[0].input_to_conv2(x) 403 | return x 404 | 405 | def layer2_0_conv2_to_output(self, x, residual): 406 | x = self.layer2[0].conv2_to_output(x, residual) 407 | x = self.layer2[1](x) 408 | x = self.layer3(x) 409 | x = self.layer4(x) 410 | x = self.avgpool(x) 411 | x = x.view(x.size(0), -1) 412 | x = self.fc(x) 413 | return x 414 | 415 | def input_to_layer2_0_conv1(self, x): 416 | x = self.conv1(x) 417 | x = self.bn1(x) 418 | x = self.relu(x) 419 | x = self.layer1(x) 420 | x = self.layer2[0].input_to_conv1(x) 421 | return x 422 | 423 | def layer2_0_conv1_to_output(self, x, residual): 424 | x = self.layer2[0].conv1_to_output(x, residual) 425 | x = self.layer2[1](x) 426 | x = self.layer3(x) 427 | x = self.layer4(x) 428 | x = self.avgpool(x) 429 | x = x.view(x.size(0), -1) 430 | x = self.fc(x) 431 | return x 432 | ############################################################################# 433 | 434 | def input_to_layer2_1_residual(self, x): 435 | x = self.conv1(x) 436 | x = self.bn1(x) 437 | x = self.relu(x) 438 | 439 | x = self.layer1(x) 440 | x = self.layer2[0](x) 441 | x = self.layer2[1].input_to_residual(x) 442 | 443 | return x 444 | 445 | def input_to_layer2_1_conv2(self, x): 446 | x = self.conv1(x) 447 | x = self.bn1(x) 448 | x = self.relu(x) 449 | x = self.layer1(x) 450 | x = self.layer2[0](x) 451 | x = self.layer2[1].input_to_conv2(x) 452 | return x 453 | 454 | def layer2_1_conv2_to_output(self, x, residual): 455 | x = self.layer2[1].conv2_to_output(x, residual) 456 | x = self.layer3(x) 457 | x = self.layer4(x) 458 | x = self.avgpool(x) 459 | x = x.view(x.size(0), -1) 460 | x = self.fc(x) 461 | return x 462 | 463 | def input_to_layer2_1_conv1(self, x): 464 | x = self.conv1(x) 465 | x = self.bn1(x) 466 | x = self.relu(x) 467 | x = self.layer1(x) 468 | x = self.layer2[0](x) 469 | x = self.layer2[1].input_to_conv1(x) 470 | return x 471 | 472 | def layer2_1_conv1_to_output(self, x, residual): 473 | x = self.layer2[1].conv1_to_output(x, residual) 474 | x = self.layer3(x) 475 | x = self.layer4(x) 476 | x = self.avgpool(x) 477 | x = x.view(x.size(0), -1) 478 | x = self.fc(x) 479 | return x 480 | 481 | ############################################################################# 482 | 483 | ############################################################################# 484 | def input_to_layer3_0_residual(self, x): 485 | x = self.conv1(x) 486 | x = self.bn1(x) 487 | x = self.relu(x) 488 | 489 | x = self.layer1(x) 490 | x = self.layer2(x) 491 | x = self.layer3[0].input_to_residual(x) 492 | 493 | return x 494 | 495 | def layer3_0_residual_to_output(self, residual, conv2): 496 | 497 | x = self.layer3[0].residual_to_output(residual,conv2) 498 | x = self.layer3[1](x) 499 | x = self.layer4(x) 500 | 501 | x = self.avgpool(x) 502 | x = x.view(x.size(0), -1) 503 | x = self.fc(x) 504 | return x 505 | 506 | def input_to_layer3_0_conv2(self, x): 507 | x = self.conv1(x) 508 | x = self.bn1(x) 509 | x = self.relu(x) 510 | x = self.layer1(x) 511 | x = self.layer2(x) 512 | x = self.layer3[0].input_to_conv2(x) 513 | return x 514 | 515 | def layer3_0_conv2_to_output(self, x, residual): 516 | x = self.layer3[0].conv2_to_output(x, residual) 517 | x = self.layer3[1](x) 518 | x = self.layer4(x) 519 | x = self.avgpool(x) 520 | x = x.view(x.size(0), -1) 521 | x = self.fc(x) 522 | return x 523 | 524 | def input_to_layer3_0_conv1(self, x): 525 | x = self.conv1(x) 526 | x = self.bn1(x) 527 | x = self.relu(x) 528 | x = self.layer1(x) 529 | x = self.layer2(x) 530 | x = self.layer3[0].input_to_conv1(x) 531 | return x 532 | 533 | def layer3_0_conv1_to_output(self, x, residual): 534 | x = self.layer3[0].conv1_to_output(x, residual) 535 | x = self.layer3[1](x) 536 | x = self.layer4(x) 537 | x = self.avgpool(x) 538 | x = x.view(x.size(0), -1) 539 | x = self.fc(x) 540 | return x 541 | ############################################################################# 542 | 543 | def input_to_layer3_1_residual(self, x): 544 | x = self.conv1(x) 545 | x = self.bn1(x) 546 | x = self.relu(x) 547 | 548 | x = self.layer1(x) 549 | x = self.layer2(x) 550 | x = self.layer3[0](x) 551 | x = self.layer3[1].input_to_residual(x) 552 | 553 | return x 554 | 555 | def input_to_layer3_1_conv2(self, x): 556 | x = self.conv1(x) 557 | x = self.bn1(x) 558 | x = self.relu(x) 559 | x = self.layer1(x) 560 | x = self.layer2(x) 561 | x = self.layer3[0](x) 562 | x = self.layer3[1].input_to_conv2(x) 563 | return x 564 | 565 | def layer3_1_conv2_to_output(self, x, residual): 566 | x = self.layer3[1].conv2_to_output(x, residual) 567 | x = self.layer4(x) 568 | x = self.avgpool(x) 569 | x = x.view(x.size(0), -1) 570 | x = self.fc(x) 571 | return x 572 | 573 | def input_to_layer3_1_conv1(self, x): 574 | x = self.conv1(x) 575 | x = self.bn1(x) 576 | x = self.relu(x) 577 | x = self.layer1(x) 578 | x = self.layer2(x) 579 | x = self.layer3[0](x) 580 | x = self.layer3[1].input_to_conv1(x) 581 | return x 582 | 583 | def layer3_1_conv1_to_output(self, x, residual): 584 | x = self.layer3[1].conv1_to_output(x, residual) 585 | x = self.layer4(x) 586 | x = self.avgpool(x) 587 | x = x.view(x.size(0), -1) 588 | x = self.fc(x) 589 | return x 590 | 591 | ############################################################################# 592 | def input_to_layer4_0_residual(self, x): 593 | x = self.conv1(x) 594 | x = self.bn1(x) 595 | x = self.relu(x) 596 | 597 | x = self.layer1(x) 598 | x = self.layer2(x) 599 | x = self.layer3(x) 600 | x = self.layer4[0].input_to_residual(x) 601 | 602 | return x 603 | 604 | def layer4_0_residual_to_output(self, residual, conv2): 605 | 606 | x = self.layer4[0].residual_to_output(residual,conv2) 607 | x = self.layer4[1](x) 608 | 609 | x = self.avgpool(x) 610 | x = x.view(x.size(0), -1) 611 | x = self.fc(x) 612 | return x 613 | 614 | def input_to_layer4_0_conv2(self, x): 615 | x = self.conv1(x) 616 | x = self.bn1(x) 617 | x = self.relu(x) 618 | x = self.layer1(x) 619 | x = self.layer2(x) 620 | x = self.layer3(x) 621 | x = self.layer4[0].input_to_conv2(x) 622 | return x 623 | 624 | def layer4_0_conv2_to_output(self, x, residual): 625 | x = self.layer4[0].conv2_to_output(x, residual) 626 | x = self.layer4[1](x) 627 | x = self.avgpool(x) 628 | x = x.view(x.size(0), -1) 629 | x = self.fc(x) 630 | return x 631 | 632 | def input_to_layer4_0_conv1(self, x): 633 | x = self.conv1(x) 634 | x = self.bn1(x) 635 | x = self.relu(x) 636 | x = self.layer1(x) 637 | x = self.layer2(x) 638 | x = self.layer3(x) 639 | x = self.layer4[0].input_to_conv1(x) 640 | return x 641 | 642 | def layer4_0_conv1_to_output(self, x, residual): 643 | x = self.layer4[0].conv1_to_output(x, residual) 644 | x = self.layer4[1](x) 645 | x = self.avgpool(x) 646 | x = x.view(x.size(0), -1) 647 | x = self.fc(x) 648 | return x 649 | ############################################################################# 650 | def input_to_layer4_1_residual(self, x): 651 | x = self.conv1(x) 652 | x = self.bn1(x) 653 | x = self.relu(x) 654 | 655 | x = self.layer1(x) 656 | x = self.layer2(x) 657 | x = self.layer3(x) 658 | x = self.layer4[0](x) 659 | x = self.layer4[1].input_to_residual(x) 660 | 661 | return x 662 | 663 | def input_to_layer4_1_conv2(self, x): 664 | x = self.conv1(x) 665 | x = self.bn1(x) 666 | x = self.relu(x) 667 | x = self.layer1(x) 668 | x = self.layer2(x) 669 | x = self.layer3(x) 670 | x = self.layer4[0](x) 671 | x = self.layer4[1].input_to_conv2(x) 672 | return x 673 | 674 | def layer4_1_conv2_to_output(self, x, residual): 675 | x = self.layer4[1].conv2_to_output(x, residual) 676 | x = self.avgpool(x) 677 | x = x.view(x.size(0), -1) 678 | x = self.fc(x) 679 | return x 680 | 681 | def input_to_layer4_1_conv1(self, x): 682 | x = self.conv1(x) 683 | x = self.bn1(x) 684 | x = self.relu(x) 685 | x = self.layer1(x) 686 | x = self.layer2(x) 687 | x = self.layer3(x) 688 | x = self.layer4[0](x) 689 | x = self.layer4[1].input_to_conv1(x) 690 | return x 691 | 692 | def layer4_1_conv1_to_output(self, x, residual): 693 | x = self.layer4[1].conv1_to_output(x, residual) 694 | x = self.avgpool(x) 695 | x = x.view(x.size(0), -1) 696 | x = self.fc(x) 697 | return x 698 | ############################################################################# 699 | 700 | def resnet18(**kwargs): 701 | return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 702 | 703 | 704 | def resnet34(**kwargs): 705 | return ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 706 | 707 | 708 | def resnet50(**kwargs): 709 | return ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 710 | 711 | 712 | def resnet101(**kwargs): 713 | return ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 714 | 715 | 716 | def resnet152(**kwargs): 717 | return ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) -------------------------------------------------------------------------------- /models/vgg.py: -------------------------------------------------------------------------------- 1 | '''VGG for CIFAR10. FC layers are removed. 2 | (c) YANG, Wei 3 | ''' 4 | import torch.nn as nn 5 | import torch.utils.model_zoo as model_zoo 6 | import math 7 | 8 | 9 | __all__ = [ 10 | 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 11 | 'vgg19_bn', 'vgg19', 12 | ] 13 | 14 | 15 | model_urls = { 16 | 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', 17 | 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', 18 | 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', 19 | 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', 20 | } 21 | 22 | 23 | class VGG(nn.Module): 24 | 25 | def __init__(self, features, num_classes=10): 26 | super(VGG, self).__init__() 27 | self.features = features 28 | self.classifier = nn.Linear(512, num_classes) 29 | self._initialize_weights() 30 | 31 | def forward(self, x): 32 | x = self.features(x) 33 | x = x.view(x.size(0), -1) 34 | x = self.classifier(x) 35 | return x 36 | 37 | def get_fm(self, x): 38 | x = self.features(x) 39 | return x 40 | 41 | def get_all_inner_activation(self, x): 42 | if len(self.features) == 31: 43 | inner_output_index = [0,2,5,7,10,12,14,17,19,21,24,26,28] 44 | #max_range = 31 45 | if len(self.features) == 21: 46 | inner_output_index = [0,3,6,8,11,13,16,18] 47 | #max_range = 21 48 | inner_output_list = [] 49 | for i in range(len(self.features)): 50 | x = self.features[i](x) 51 | if i in inner_output_index: 52 | inner_output_list.append(x) 53 | x = x.view(x.size(0), -1) 54 | x = self.classifier(x) 55 | return x,inner_output_list 56 | 57 | def _initialize_weights(self): 58 | for m in self.modules(): 59 | if isinstance(m, nn.Conv2d): 60 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 61 | m.weight.data.normal_(0, math.sqrt(2. / n)) 62 | if m.bias is not None: 63 | m.bias.data.zero_() 64 | elif isinstance(m, nn.BatchNorm2d): 65 | m.weight.data.fill_(1) 66 | m.bias.data.zero_() 67 | elif isinstance(m, nn.Linear): 68 | n = m.weight.size(1) 69 | m.weight.data.normal_(0, 0.01) 70 | m.bias.data.zero_() 71 | 72 | def make_layers(cfg, batch_norm=False): 73 | layers = [] 74 | in_channels = 3 75 | for v in cfg: 76 | if v == 'M': 77 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 78 | else: 79 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 80 | if batch_norm: 81 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=False)] 82 | else: 83 | layers += [conv2d, nn.ReLU(inplace=False)] 84 | in_channels = v 85 | return nn.Sequential(*layers) 86 | 87 | 88 | cfg = { 89 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 90 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 91 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 92 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 93 | } 94 | 95 | 96 | def vgg11(**kwargs): 97 | """VGG 11-layer model (configuration "A") 98 | Args: 99 | pretrained (bool): If True, returns a model pre-trained on ImageNet 100 | """ 101 | model = VGG(make_layers(cfg['A']), **kwargs) 102 | return model 103 | 104 | 105 | def vgg11_bn(**kwargs): 106 | """VGG 11-layer model (configuration "A") with batch normalization""" 107 | model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs) 108 | return model 109 | 110 | 111 | def vgg13(**kwargs): 112 | """VGG 13-layer model (configuration "B") 113 | Args: 114 | pretrained (bool): If True, returns a model pre-trained on ImageNet 115 | """ 116 | model = VGG(make_layers(cfg['B']), **kwargs) 117 | return model 118 | 119 | 120 | def vgg13_bn(**kwargs): 121 | """VGG 13-layer model (configuration "B") with batch normalization""" 122 | model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs) 123 | return model 124 | 125 | 126 | def vgg16(**kwargs): 127 | """VGG 16-layer model (configuration "D") 128 | Args: 129 | pretrained (bool): If True, returns a model pre-trained on ImageNet 130 | """ 131 | model = VGG(make_layers(cfg['D']), **kwargs) 132 | return model 133 | 134 | 135 | def vgg16_bn(**kwargs): 136 | """VGG 16-layer model (configuration "D") with batch normalization""" 137 | model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs) 138 | return model 139 | 140 | 141 | def vgg19(**kwargs): 142 | """VGG 19-layer model (configuration "E") 143 | Args: 144 | pretrained (bool): If True, returns a model pre-trained on ImageNet 145 | """ 146 | model = VGG(make_layers(cfg['E']), **kwargs) 147 | return model 148 | 149 | 150 | def vgg19_bn(**kwargs): 151 | """VGG 19-layer model (configuration 'E') with batch normalization""" 152 | model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs) 153 | return model 154 | -------------------------------------------------------------------------------- /models/vgg_trojai.py: -------------------------------------------------------------------------------- 1 | '''VGG for CIFAR10. FC layers are removed. 2 | (c) YANG, Wei 3 | ''' 4 | import torch.nn as nn 5 | import torch.utils.model_zoo as model_zoo 6 | import math 7 | 8 | 9 | __all__ = [ 10 | 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 11 | 'vgg19_bn', 'vgg19', 12 | ] 13 | 14 | 15 | model_urls = { 16 | 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', 17 | 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', 18 | 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', 19 | 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', 20 | } 21 | 22 | 23 | class VGG(nn.Module): 24 | 25 | def __init__(self, features, num_classes=10): 26 | super(VGG, self).__init__() 27 | self.features = features 28 | #self.classifier = nn.Linear(512, num_classes) 29 | self.avgpool = nn.AdaptiveAvgPool2d(output_size=(7, 7)) 30 | self.classifier = nn.Sequential( 31 | nn.Linear(512 * 7 * 7, 4096), 32 | nn.ReLU(True), 33 | nn.Dropout(), 34 | nn.Linear(4096, 4096), 35 | nn.ReLU(True), 36 | nn.Dropout(), 37 | nn.Linear(4096, num_classes), 38 | ) 39 | self._initialize_weights() 40 | 41 | def forward(self, x): 42 | x = self.features(x) 43 | x = self.avgpool(x) 44 | x = x.view(x.size(0), -1) 45 | x = self.classifier(x) 46 | return x 47 | 48 | def get_fm(self, x): 49 | x = self.features(x) 50 | return x 51 | 52 | def get_all_inner_activation(self, x): 53 | if len(self.features) == 31: 54 | inner_output_index = [0,2,5,7,10,12,14,17,19,21,24,26,28] 55 | #max_range = 31 56 | if len(self.features) == 21: 57 | inner_output_index = [0,3,6,8,11,13,16,18] 58 | #max_range = 21 59 | inner_output_list = [] 60 | for i in range(len(self.features)): 61 | x = self.features[i](x) 62 | if i in inner_output_index: 63 | inner_output_list.append(x) 64 | x = x.view(x.size(0), -1) 65 | x = self.classifier(x) 66 | return x,inner_output_list 67 | 68 | def _initialize_weights(self): 69 | for m in self.modules(): 70 | if isinstance(m, nn.Conv2d): 71 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 72 | m.weight.data.normal_(0, math.sqrt(2. / n)) 73 | if m.bias is not None: 74 | m.bias.data.zero_() 75 | elif isinstance(m, nn.BatchNorm2d): 76 | m.weight.data.fill_(1) 77 | m.bias.data.zero_() 78 | elif isinstance(m, nn.Linear): 79 | n = m.weight.size(1) 80 | m.weight.data.normal_(0, 0.01) 81 | m.bias.data.zero_() 82 | 83 | 84 | def make_layers(cfg, batch_norm=False): 85 | layers = [] 86 | in_channels = 3 87 | for v in cfg: 88 | if v == 'M': 89 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 90 | else: 91 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 92 | if batch_norm: 93 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 94 | else: 95 | layers += [conv2d, nn.ReLU(inplace=True)] 96 | in_channels = v 97 | return nn.Sequential(*layers) 98 | 99 | 100 | cfg = { 101 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 102 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 103 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 104 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 105 | } 106 | 107 | 108 | def vgg11(**kwargs): 109 | """VGG 11-layer model (configuration "A") 110 | Args: 111 | pretrained (bool): If True, returns a model pre-trained on ImageNet 112 | """ 113 | model = VGG(make_layers(cfg['A']), **kwargs) 114 | return model 115 | 116 | 117 | def vgg11_bn(**kwargs): 118 | """VGG 11-layer model (configuration "A") with batch normalization""" 119 | model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs) 120 | return model 121 | 122 | 123 | def vgg13(**kwargs): 124 | """VGG 13-layer model (configuration "B") 125 | Args: 126 | pretrained (bool): If True, returns a model pre-trained on ImageNet 127 | """ 128 | model = VGG(make_layers(cfg['B']), **kwargs) 129 | return model 130 | 131 | 132 | def vgg13_bn(**kwargs): 133 | """VGG 13-layer model (configuration "B") with batch normalization""" 134 | model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs) 135 | return model 136 | 137 | 138 | def vgg16(**kwargs): 139 | """VGG 16-layer model (configuration "D") 140 | Args: 141 | pretrained (bool): If True, returns a model pre-trained on ImageNet 142 | """ 143 | model = VGG(make_layers(cfg['D']), **kwargs) 144 | return model 145 | 146 | 147 | def vgg16_bn(**kwargs): 148 | """VGG 16-layer model (configuration "D") with batch normalization""" 149 | model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs) 150 | return model 151 | 152 | 153 | def vgg19(**kwargs): 154 | """VGG 19-layer model (configuration "E") 155 | Args: 156 | pretrained (bool): If True, returns a model pre-trained on ImageNet 157 | """ 158 | model = VGG(make_layers(cfg['E']), **kwargs) 159 | return model 160 | 161 | 162 | def vgg19_bn(**kwargs): 163 | """VGG 19-layer model (configuration 'E') with batch normalization""" 164 | model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs) 165 | return model 166 | -------------------------------------------------------------------------------- /none/compromised_detection.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | import jenkspy 4 | from torch.autograd import Variable 5 | import torch.nn.functional as F 6 | 7 | def analyze_neuros(model,arch,max_reset_fraction,lamda_l,lamda_h, \ 8 | num_classes,num_for_detect_biased,trainloader_no_shuffle): 9 | model.eval() 10 | selected_neuros = [] 11 | index = 0 12 | #loader = trainloader 13 | loader = trainloader_no_shuffle 14 | #max_reset_fraction = args.reset_max_per_layer_fraction 15 | 16 | if arch == "nin": 17 | conv_list = ["Conv_0_","Conv_2_","Conv_4_","Conv_8_","Conv_10_","Conv_12_","Conv_16_","Conv_18_"] 18 | channel_num_list = [192,160,96,192,192,192,192,192] 19 | elif arch == "vgg16": 20 | conv_list = ["Conv_0_","Conv_2_","Conv_5_","Conv_7_","Conv_10_","Conv_12_","Conv_14_","Conv_17_","Conv_19_","Conv_21_","Conv_24_","Conv_26_","Conv_28_"] 21 | channel_num_list = [64,64,128,128,256,256,256,512,512,512,512,512,512] 22 | elif arch in ["resnet18","resnet18_leakyrelu","resnet18_elu","resnet18_prelu","resnet18_tanhshrink","resnet18_softplus"]: 23 | conv_list = ["Layer2_0_Conv1_","Layer2_0_Downsample_","Layer2_0_Conv2_","Layer2_1_Conv1_","Layer2_1_Conv2_","Layer3_0_Conv1_","Layer3_0_Downsample_","Layer3_0_Conv2_","Layer3_1_Conv1_","Layer3_1_Conv2_","Layer4_0_Conv1_","Layer4_0_Downsample_","Layer4_0_Conv2_","Layer4_1_Conv1_","Layer4_1_Conv2_"] 24 | channel_num_list = [128,128,128,128,128,256,256,256,256,256,512,512,512,512,512] 25 | 26 | imgs_high_activation_times = np.zeros([len(loader.dataset)]) 27 | imgs_high_activation_times_counter = np.zeros([len(loader.dataset)]) 28 | conv_activation_all_list = [] 29 | 30 | for i in range(len(conv_list)): 31 | channel_num = channel_num_list[i] 32 | conv_activation_all_list.append(np.zeros([len(loader.dataset),channel_num])) 33 | 34 | output_all = np.zeros([len(loader.dataset),num_classes]) 35 | 36 | counter = 0 37 | 38 | if arch in ["resnet18","resnet18_leakyrelu","resnet18_elu","resnet18_prelu","resnet18_tanhshrink","resnet18_softplus","efficientnet","mobilenetv2","densenet121"]: 39 | for data, target in loader: 40 | data, target = Variable(data.cuda()), Variable(target.cuda()) 41 | output = model(data) 42 | pred = F.softmax(output) 43 | batch_size = data.shape[0] 44 | output_all[counter:counter+batch_size] = pred.cpu().detach().numpy() 45 | for i in range(len(conv_list)): 46 | conv_name = conv_list[i] 47 | if "FC" in conv_name: 48 | conv_activation_all_list[i][counter:counter+batch_size] = model.inter_feature[conv_name].cpu().detach().numpy() 49 | else: 50 | conv_activation_all_list[i][counter:counter+batch_size] = model.inter_feature[conv_name].max(-1).values.max(-1).values.cpu().detach().numpy() 51 | counter = counter+batch_size 52 | print(counter) 53 | 54 | else: 55 | target_count = [0]*10 56 | for data, target in loader: 57 | data, target = Variable(data.cuda()), Variable(target.cuda()) 58 | output, inner_output_list = model.get_all_inner_activation(data) 59 | pred = F.softmax(output) 60 | for i in range(target.shape[0]): 61 | target_count[target[i].item()] = target_count[target[i].item()] + 1 62 | batch_size = data.shape[0] 63 | output_all[counter:counter+batch_size] = pred.cpu().detach().numpy() 64 | for i in range(len(conv_list)): 65 | conv_name = conv_list[i] 66 | if "FC" in conv_name: 67 | conv_activation_all_list[i][counter:counter+batch_size] = inner_output_list[i].cpu().detach().numpy() 68 | else: 69 | conv_activation_all_list[i][counter:counter+batch_size] = inner_output_list[i].max(-1).values.max(-1).values.cpu().detach().numpy() 70 | counter = counter+batch_size 71 | print(counter) 72 | print(target_count) 73 | 74 | diff_class_channel_numpy_list = [] 75 | for i in range(len(conv_list)): 76 | channel_num = channel_num_list[i] 77 | diff_class_channel_numpy_list.append(np.zeros([channel_num,num_classes])) 78 | 79 | for i in range(len(conv_list)): 80 | conv_name = conv_list[i] 81 | channel_num = channel_num_list[i] 82 | print(conv_name) 83 | print(channel_num) 84 | for j in range(channel_num): 85 | for k in range(num_classes): 86 | strong_output_indexs = np.where(output_all[:,k]>lamda_h) 87 | weak_output_indexs = np.where(output_all[:,k]lamda_h) 114 | weak_output_indexs = np.where(output_all[:,k]=breaks[1])[0].tolist() 120 | print("number of detected Trojan samples:",len(poison_sample_index)) 121 | return selected_neuros, poison_sample_index 122 | -------------------------------------------------------------------------------- /none/reset_neurons.py: -------------------------------------------------------------------------------- 1 | 2 | def reset(model,arch,selected_neuros,freeze = False): 3 | model.eval() 4 | channel_indices=[] 5 | 6 | if arch == "vgg16": 7 | 8 | conv_list = ["Conv_0_","Conv_2_","Conv_5_","Conv_7_","Conv_10_","Conv_12_","Conv_14_","Conv_17_","Conv_19_","Conv_21_","Conv_24_","Conv_26_","Conv_28_"] 9 | layer_list = [0,2,5,7,10,12,14,17,19,21,24,26,28] 10 | 11 | for j in range(len(layer_list)): 12 | print("reset Conv "+str(layer_list[j])+":") 13 | channel_indices=[] 14 | for line in selected_neuros: 15 | layer = line.split("_")[1] 16 | neuro = line.split("_")[2] 17 | if layer == str(layer_list[j]): 18 | channel_indices.append(int(neuro)) 19 | print(int(neuro)) 20 | for name, param in model.named_parameters(): 21 | for i in range(0,len(channel_indices)): 22 | channel = channel_indices[i] 23 | if name == "features."+str(layer_list[j])+".weight": 24 | param.data[channel,:,:,:] = 0 25 | if name == "features."+str(layer_list[j])+".bias": 26 | param.data[channel] = 0 27 | 28 | if (j == (len(layer_list)-1)): 29 | if name == "classifier.weight": 30 | param.data[:,channel] = 0 31 | else: 32 | if name == "features."+str(layer_list[j+1])+".weight": 33 | param.data[:,channel,:,:] = 0 34 | 35 | elif arch == "nin": 36 | 37 | print("reset layer 0:") 38 | channel_indices=[] 39 | for line in selected_neuros: 40 | layer = line.split("_")[1] 41 | neuro = line.split("_")[2] 42 | if layer == "0": 43 | channel_indices.append(int(neuro)) 44 | print(int(neuro)) 45 | for name, param in model.named_parameters(): 46 | for i in range(0,len(channel_indices)): 47 | channel = channel_indices[i] 48 | if name == "classifier.0.weight": 49 | param.data[channel,:,:,:] = 0 50 | if name == "classifier.0.bias": 51 | param.data[channel] = 0 52 | if name == "classifier.2.weight": 53 | param.data[:,channel,:,:] = 0 54 | 55 | if freeze: 56 | for i in range(0,len(channel_indices)): 57 | channel = channel_indices[i] 58 | if name == "classifier.0.weight": 59 | param.data[channel,:,:,:] = 0 60 | param.requires_grad = False 61 | if name == "classifier.0.bias": 62 | param.data[channel] = 0 63 | param.requires_grad = False 64 | if name == "classifier.2.weight": 65 | param.data[:,channel,:,:] = 0 66 | param.requires_grad = False 67 | 68 | print("reset layer 2:") 69 | channel_indices=[] 70 | for line in selected_neuros: 71 | layer = line.split("_")[1] 72 | neuro = line.split("_")[2] 73 | if layer == "2": 74 | channel_indices.append(int(neuro)) 75 | print(int(neuro)) 76 | for name, param in model.named_parameters(): 77 | for i in range(0,len(channel_indices)): 78 | channel = channel_indices[i] 79 | if name == "classifier.2.weight": 80 | param.data[channel,:,:,:] = 0 81 | if name == "classifier.2.bias": 82 | param.data[channel] = 0 83 | if name == "classifier.4.weight": 84 | param.data[:,channel,:,:] = 0 85 | 86 | if freeze: 87 | for i in range(0,len(channel_indices)): 88 | channel = channel_indices[i] 89 | if name == "classifier.0.weight": 90 | param.data[channel,:,:,:] = 0 91 | param.requires_grad = False 92 | if name == "classifier.0.bias": 93 | param.data[channel] = 0 94 | param.requires_grad = False 95 | if name == "classifier.2.weight": 96 | param.data[:,channel,:,:] = 0 97 | param.requires_grad = False 98 | 99 | print("reset layer 4:") 100 | channel_indices=[] 101 | for line in selected_neuros: 102 | layer = line.split("_")[1] 103 | neuro = line.split("_")[2] 104 | if layer == "4": 105 | channel_indices.append(int(neuro)) 106 | print(int(neuro)) 107 | for name, param in model.named_parameters(): 108 | for i in range(0,len(channel_indices)): 109 | channel = channel_indices[i] 110 | if name == "classifier.4.weight": 111 | param.data[channel,:,:,:] = 0 112 | if name == "classifier.4.bias": 113 | param.data[channel] = 0 114 | if name == "classifier.8.weight": 115 | param.data[:,channel,:,:] = 0 116 | 117 | if freeze: 118 | for i in range(0,len(channel_indices)): 119 | channel = channel_indices[i] 120 | if name == "classifier.4.weight": 121 | param.data[channel,:,:,:] = 0 122 | param.requires_grad = False 123 | if name == "classifier.4.bias": 124 | param.data[channel] = 0 125 | param.requires_grad = False 126 | if name == "classifier.8.weight": 127 | param.data[:,channel,:,:] = 0 128 | param.requires_grad = False 129 | 130 | print("reset layer 8:") 131 | channel_indices=[] 132 | for line in selected_neuros: 133 | layer = line.split("_")[1] 134 | neuro = line.split("_")[2] 135 | if layer == "8": 136 | channel_indices.append(int(neuro)) 137 | print(int(neuro)) 138 | for name, param in model.named_parameters(): 139 | for i in range(0,len(channel_indices)): 140 | channel = channel_indices[i] 141 | if name == "classifier.8.weight": 142 | param.data[channel,:,:,:] = 0 143 | if name == "classifier.8.bias": 144 | param.data[channel] = 0 145 | if name == "classifier.10.weight": 146 | param.data[:,channel,:,:] = 0 147 | 148 | if freeze: 149 | for i in range(0,len(channel_indices)): 150 | channel = channel_indices[i] 151 | if name == "classifier.8.weight": 152 | param.data[channel,:,:,:] = 0 153 | param.requires_grad = False 154 | if name == "classifier.8.bias": 155 | param.data[channel] = 0 156 | param.requires_grad = False 157 | if name == "classifier.10.weight": 158 | param.data[:,channel,:,:] = 0 159 | param.requires_grad = False 160 | 161 | print("reset layer 10:") 162 | channel_indices=[] 163 | for line in selected_neuros: 164 | layer = line.split("_")[1] 165 | neuro = line.split("_")[2] 166 | if layer == "10": 167 | channel_indices.append(int(neuro)) 168 | print(int(neuro)) 169 | for name, param in model.named_parameters(): 170 | for i in range(0,len(channel_indices)): 171 | channel = channel_indices[i] 172 | if name == "classifier.10.weight": 173 | param.data[channel,:,:,:] = 0 174 | if name == "classifier.10.bias": 175 | param.data[channel] = 0 176 | if name == "classifier.12.weight": 177 | param.data[:,channel,:,:] = 0 178 | 179 | if freeze: 180 | for i in range(0,len(channel_indices)): 181 | channel = channel_indices[i] 182 | if name == "classifier.10.weight": 183 | param.data[channel,:,:,:] = 0 184 | param.requires_grad = False 185 | if name == "classifier.10.bias": 186 | param.data[channel] = 0 187 | param.requires_grad = False 188 | if name == "classifier.12.weight": 189 | param.data[:,channel,:,:] = 0 190 | param.requires_grad = False 191 | 192 | print("reset layer 12:") 193 | channel_indices=[] 194 | for line in selected_neuros: 195 | layer = line.split("_")[1] 196 | neuro = line.split("_")[2] 197 | if layer == "12": 198 | channel_indices.append(int(neuro)) 199 | print(int(neuro)) 200 | for name, param in model.named_parameters(): 201 | for i in range(0,len(channel_indices)): 202 | channel = channel_indices[i] 203 | if name == "classifier.12.weight": 204 | param.data[channel,:,:,:] = 0 205 | if name == "classifier.12.bias": 206 | param.data[channel] = 0 207 | if name == "classifier.16.weight": 208 | param.data[:,channel,:,:] = 0 209 | 210 | if freeze: 211 | for i in range(0,len(channel_indices)): 212 | channel = channel_indices[i] 213 | if name == "classifier.12.weight": 214 | param.data[channel,:,:,:] = 0 215 | param.requires_grad = False 216 | if name == "classifier.12.bias": 217 | param.data[channel] = 0 218 | param.requires_grad = False 219 | if name == "classifier.16.weight": 220 | param.data[:,channel,:,:] = 0 221 | param.requires_grad = False 222 | 223 | print("reset layer 16:") 224 | channel_indices=[] 225 | for line in selected_neuros: 226 | layer = line.split("_")[1] 227 | neuro = line.split("_")[2] 228 | if layer == "16": 229 | channel_indices.append(int(neuro)) 230 | print(int(neuro)) 231 | for name, param in model.named_parameters(): 232 | 233 | for i in range(0,len(channel_indices)): 234 | channel = channel_indices[i] 235 | if name == "classifier.16.weight": 236 | param.data[channel,:,:,:] = 0 237 | if name == "classifier.16.bias": 238 | param.data[channel] = 0 239 | if name == "classifier.18.weight": 240 | param.data[:,channel,:,:] = 0 241 | 242 | if freeze: 243 | for i in range(0,len(channel_indices)): 244 | channel = channel_indices[i] 245 | if name == "classifier.16.weight": 246 | param.data[channel,:,:,:] = 0 247 | param.requires_grad = False 248 | if name == "classifier.16.bias": 249 | param.data[channel] = 0 250 | param.requires_grad = False 251 | if name == "classifier.18.weight": 252 | param.data[:,channel,:,:] = 0 253 | param.requires_grad = False 254 | 255 | print("reset layer 18:") 256 | channel_indices=[] 257 | for line in selected_neuros: 258 | layer = line.split("_")[1] 259 | neuro = line.split("_")[2] 260 | if layer == "18": 261 | channel_indices.append(int(neuro)) 262 | print(int(neuro)) 263 | 264 | for name, param in model.named_parameters(): 265 | for i in range(0,len(channel_indices)): 266 | channel = channel_indices[i] 267 | if name == "classifier.18.weight": 268 | param.data[channel,:,:,:] = 0 269 | if name == "classifier.18.bias": 270 | param.data[channel] = 0 271 | if name == "classifier.20.weight": 272 | param.data[:,channel,:,:] = 0 273 | 274 | if freeze: 275 | for i in range(0,len(channel_indices)): 276 | channel = channel_indices[i] 277 | if name == "classifier.18.weight": 278 | param.data[channel,:,:,:] = 0 279 | param.requires_grad = False 280 | if name == "classifier.18.bias": 281 | param.data[channel] = 0 282 | param.requires_grad = False 283 | if name == "classifier.20.weight": 284 | param.data[:,channel,:,:] = 0 285 | param.requires_grad = False 286 | 287 | elif arch in ["resnet18","resnet18_leakyrelu","resnet18_elu","resnet18_prelu","resnet18_tanhshrink","resnet18_softplus"]: 288 | 289 | print("reset Conv1:") 290 | channel_indices=[] 291 | for line in selected_neuros: 292 | layer = line.split("_")[0] 293 | sublayer = line.split("_")[1] 294 | conv = line.split("_")[2] 295 | neuro = line.split("_")[3] 296 | if (layer == "Conv1") and (sublayer == "Conv1") and (conv == "Conv1"): 297 | channel_indices.append(int(neuro)) 298 | print(int(neuro)) 299 | for name, param in model.named_parameters(): 300 | for i in range(0,len(channel_indices)): 301 | channel = channel_indices[i] 302 | if name == "conv1.weight": 303 | param.data[channel,:,:,:] = 0 304 | if name == "bn1.weight": 305 | param.data[channel] = 0 306 | if name == "bn1.bias": 307 | param.data[channel] = 0 308 | if name == "layer1.0.conv1.weight": 309 | param.data[:,channel,:,:] = 0 310 | 311 | print("reset Layer1.0.Conv1:") 312 | channel_indices=[] 313 | for line in selected_neuros: 314 | layer = line.split("_")[0] 315 | sublayer = line.split("_")[1] 316 | conv = line.split("_")[2] 317 | neuro = line.split("_")[3] 318 | 319 | if (layer == "Layer1") and (sublayer == "0") and (conv == "Conv1"): 320 | channel_indices.append(int(neuro)) 321 | print(int(neuro)) 322 | for name, param in model.named_parameters(): 323 | for i in range(0,len(channel_indices)): 324 | channel = channel_indices[i] 325 | if name == "layer1.0.conv1.weight": 326 | param.data[channel,:,:,:] = 0 327 | if name == "layer1.0.bn1.weight": 328 | param.data[channel] = 0 329 | if name == "layer1.0.bn1.bias": 330 | param.data[channel] = 0 331 | if name == "layer1.0.conv2.weight": 332 | param.data[:,channel,:,:] = 0 333 | 334 | print("reset Layer1.0.Conv2:") 335 | channel_indices=[] 336 | for line in selected_neuros: 337 | layer = line.split("_")[0] 338 | sublayer = line.split("_")[1] 339 | conv = line.split("_")[2] 340 | neuro = line.split("_")[3] 341 | 342 | if (layer == "Layer1") and (sublayer == "0") and (conv == "Conv2"): 343 | channel_indices.append(int(neuro)) 344 | print(int(neuro)) 345 | for name, param in model.named_parameters(): 346 | for i in range(0,len(channel_indices)): 347 | channel = channel_indices[i] 348 | if name == "layer1.0.conv2.weight": 349 | param.data[channel,:,:,:] = 0 350 | if name == "layer1.0.bn2.weight": 351 | param.data[channel] = 0 352 | if name == "layer1.0.bn2.bias": 353 | param.data[channel] = 0 354 | if name == "layer1.1.conv1.weight": 355 | param.data[:,channel,:,:] = 0 356 | 357 | print("reset Layer1.1.Conv1:") 358 | channel_indices=[] 359 | for line in selected_neuros: 360 | layer = line.split("_")[0] 361 | sublayer = line.split("_")[1] 362 | conv = line.split("_")[2] 363 | neuro = line.split("_")[3] 364 | 365 | if (layer == "Layer1") and (sublayer == "1") and (conv == "Conv1"): 366 | channel_indices.append(int(neuro)) 367 | print(int(neuro)) 368 | for name, param in model.named_parameters(): 369 | for i in range(0,len(channel_indices)): 370 | channel = channel_indices[i] 371 | if name == "layer1.1.conv1.weight": 372 | param.data[channel,:,:,:] = 0 373 | if name == "layer1.1.bn1.weight": 374 | param.data[channel] = 0 375 | if name == "layer1.1.bn1.bias": 376 | param.data[channel] = 0 377 | 378 | if name == "layer1.1.conv2.weight": 379 | param.data[:,channel,:,:] = 0 380 | 381 | print("reset Layer1.1.Conv2:") 382 | channel_indices=[] 383 | for line in selected_neuros: 384 | layer = line.split("_")[0] 385 | sublayer = line.split("_")[1] 386 | conv = line.split("_")[2] 387 | neuro = line.split("_")[3] 388 | 389 | if (layer == "Layer1") and (sublayer == "1") and (conv == "Conv2"): 390 | channel_indices.append(int(neuro)) 391 | print(int(neuro)) 392 | 393 | for name, param in model.named_parameters(): 394 | for i in range(0,len(channel_indices)): 395 | channel = channel_indices[i] 396 | if name == "layer1.1.conv2.weight": 397 | param.data[channel,:,:,:] = 0 398 | if name == "layer1.1.bn2.weight": 399 | param.data[channel] = 0 400 | if name == "layer1.1.bn2.bias": 401 | param.data[channel] = 0 402 | ########################################## 403 | if name == "layer2.0.conv1.weight": 404 | param.data[:,channel,:,:] = 0 405 | 406 | print("reset Layer2.0.Conv1:") 407 | channel_indices=[] 408 | for line in selected_neuros: 409 | layer = line.split("_")[0] 410 | sublayer = line.split("_")[1] 411 | conv = line.split("_")[2] 412 | neuro = line.split("_")[3] 413 | 414 | if (layer == "Layer2") and (sublayer == "0") and (conv == "Conv1"): 415 | channel_indices.append(int(neuro)) 416 | print(int(neuro)) 417 | for name, param in model.named_parameters(): 418 | for i in range(0,len(channel_indices)): 419 | channel = channel_indices[i] 420 | if name == "layer2.0.conv1.weight": 421 | param.data[channel,:,:,:] = 0 422 | if name == "layer2.0.bn1.weight": 423 | param.data[channel] = 0 424 | if name == "layer2.0.bn1.bias": 425 | param.data[channel] = 0 426 | if name == "layer2.0.conv2.weight": 427 | param.data[:,channel,:,:] = 0 428 | 429 | 430 | print("reset Layer2.0.Downsample:") 431 | channel_indices=[] 432 | for line in selected_neuros: 433 | layer = line.split("_")[0] 434 | sublayer = line.split("_")[1] 435 | conv = line.split("_")[2] 436 | neuro = line.split("_")[3] 437 | 438 | if (layer == "Layer2") and (sublayer == "0") and (conv == "Downsample"): 439 | channel_indices.append(int(neuro)) 440 | print(int(neuro)) 441 | for name, param in model.named_parameters(): 442 | for i in range(0,len(channel_indices)): 443 | channel = channel_indices[i] 444 | if name == "layer2.0.downsample.0.weight": 445 | param.data[channel,:,:,:] = 0 446 | if name == "layer2.0.downsample.1.weight": 447 | param.data[channel] = 0 448 | if name == "layer2.0.downsample.1.bias": 449 | param.data[channel] = 0 450 | 451 | print("reset Layer2.0.Conv2:") 452 | channel_indices=[] 453 | for line in selected_neuros: 454 | layer = line.split("_")[0] 455 | sublayer = line.split("_")[1] 456 | conv = line.split("_")[2] 457 | neuro = line.split("_")[3] 458 | 459 | if (layer == "Layer2") and (sublayer == "0") and (conv == "Conv2"): 460 | channel_indices.append(int(neuro)) 461 | print(int(neuro)) 462 | for name, param in model.named_parameters(): 463 | for i in range(0,len(channel_indices)): 464 | channel = channel_indices[i] 465 | if name == "layer2.0.conv2.weight": 466 | param.data[channel,:,:,:] = 0 467 | if name == "layer2.0.bn2.weight": 468 | param.data[channel] = 0 469 | if name == "layer2.0.bn2.bias": 470 | param.data[channel] = 0 471 | if name == "layer2.1.conv1.weight": 472 | param.data[:,channel,:,:] = 0 473 | 474 | print("reset Layer2.1.Conv1:") 475 | channel_indices=[] 476 | for line in selected_neuros: 477 | layer = line.split("_")[0] 478 | sublayer = line.split("_")[1] 479 | conv = line.split("_")[2] 480 | neuro = line.split("_")[3] 481 | 482 | if (layer == "Layer2") and (sublayer == "1") and (conv == "Conv1"): 483 | channel_indices.append(int(neuro)) 484 | print(int(neuro)) 485 | for name, param in model.named_parameters(): 486 | for i in range(0,len(channel_indices)): 487 | channel = channel_indices[i] 488 | if name == "layer2.1.conv1.weight": 489 | param.data[channel,:,:,:] = 0 490 | if name == "layer2.1.bn1.weight": 491 | param.data[channel] = 0 492 | if name == "layer2.1.bn1.bias": 493 | param.data[channel] = 0 494 | 495 | if name == "layer2.1.conv2.weight": 496 | param.data[:,channel,:,:] = 0 497 | 498 | print("reset Layer2.1.Conv2:") 499 | channel_indices=[] 500 | for line in selected_neuros: 501 | layer = line.split("_")[0] 502 | sublayer = line.split("_")[1] 503 | conv = line.split("_")[2] 504 | neuro = line.split("_")[3] 505 | 506 | if (layer == "Layer2") and (sublayer == "1") and (conv == "Conv2"): 507 | channel_indices.append(int(neuro)) 508 | print(int(neuro)) 509 | 510 | for name, param in model.named_parameters(): 511 | for i in range(0,len(channel_indices)): 512 | channel = channel_indices[i] 513 | if name == "layer2.1.conv2.weight": 514 | param.data[channel,:,:,:] = 0 515 | if name == "layer2.1.bn2.weight": 516 | param.data[channel] = 0 517 | if name == "layer2.1.bn2.bias": 518 | param.data[channel] = 0 519 | ########################################## 520 | if name == "layer3.0.conv1.weight": 521 | param.data[:,channel,:,:] = 0 522 | 523 | print("reset Layer3.0.Conv1:") 524 | channel_indices=[] 525 | for line in selected_neuros: 526 | layer = line.split("_")[0] 527 | sublayer = line.split("_")[1] 528 | conv = line.split("_")[2] 529 | neuro = line.split("_")[3] 530 | 531 | if (layer == "Layer3") and (sublayer == "0") and (conv == "Conv1"): 532 | channel_indices.append(int(neuro)) 533 | for name, param in model.named_parameters(): 534 | for i in range(0,len(channel_indices)): 535 | channel = channel_indices[i] 536 | if name == "layer3.0.conv1.weight": 537 | param.data[channel,:,:,:] = 0 538 | if name == "layer3.0.bn1.weight": 539 | param.data[channel] = 0 540 | if name == "layer3.0.bn1.bias": 541 | param.data[channel] = 0 542 | 543 | if name == "layer3.0.conv2.weight": 544 | param.data[:,channel,:,:] = 0 545 | 546 | print("reset Layer3.0.Downsample:") 547 | channel_indices=[] 548 | for line in selected_neuros: 549 | layer = line.split("_")[0] 550 | sublayer = line.split("_")[1] 551 | conv = line.split("_")[2] 552 | neuro = line.split("_")[3] 553 | 554 | if (layer == "Layer3") and (sublayer == "0") and (conv == "Downsample"): 555 | channel_indices.append(int(neuro)) 556 | for name, param in model.named_parameters(): 557 | for i in range(0,len(channel_indices)): 558 | channel = channel_indices[i] 559 | if name == "layer3.0.downsample.0.weight": 560 | param.data[channel,:,:,:] = 0 561 | if name == "layer3.0.downsample.1.weight": 562 | param.data[channel] = 0 563 | if name == "layer3.0.downsample.1.bias": 564 | param.data[channel] = 0 565 | 566 | print("reset Layer3.0.Conv2:") 567 | channel_indices=[] 568 | for line in selected_neuros: 569 | layer = line.split("_")[0] 570 | sublayer = line.split("_")[1] 571 | conv = line.split("_")[2] 572 | neuro = line.split("_")[3] 573 | 574 | if (layer == "Layer3") and (sublayer == "0") and (conv == "Conv2"): 575 | channel_indices.append(int(neuro)) 576 | for name, param in model.named_parameters(): 577 | for i in range(0,len(channel_indices)): 578 | channel = channel_indices[i] 579 | if name == "layer3.0.conv2.weight": 580 | param.data[channel,:,:,:] = 0 581 | if name == "layer3.0.bn2.weight": 582 | param.data[channel] = 0 583 | if name == "layer3.0.bn2.bias": 584 | param.data[channel] = 0 585 | 586 | if name == "layer3.1.conv1.weight": 587 | param.data[:,channel,:,:] = 0 588 | 589 | print("reset Layer3.1.Conv1:") 590 | channel_indices=[] 591 | for line in selected_neuros: 592 | layer = line.split("_")[0] 593 | sublayer = line.split("_")[1] 594 | conv = line.split("_")[2] 595 | neuro = line.split("_")[3] 596 | 597 | 598 | if (layer == "Layer3") and (sublayer == "1") and (conv == "Conv1"): 599 | channel_indices.append(int(neuro)) 600 | for name, param in model.named_parameters(): 601 | for i in range(0,len(channel_indices)): 602 | channel = channel_indices[i] 603 | if name == "layer3.1.conv1.weight": 604 | param.data[channel,:,:,:] = 0 605 | if name == "layer3.1.bn1.weight": 606 | param.data[channel] = 0 607 | if name == "layer3.1.bn1.bias": 608 | param.data[channel] = 0 609 | 610 | if name == "layer3.1.conv2.weight": 611 | param.data[:,channel,:,:] = 0 612 | 613 | print("reset Layer3.1.Conv2:") 614 | channel_indices=[] 615 | for line in selected_neuros: 616 | layer = line.split("_")[0] 617 | sublayer = line.split("_")[1] 618 | conv = line.split("_")[2] 619 | neuro = line.split("_")[3] 620 | 621 | if (layer == "Layer3") and (sublayer == "1") and (conv == "Conv2"): 622 | channel_indices.append(int(neuro)) 623 | for name, param in model.named_parameters(): 624 | for i in range(0,len(channel_indices)): 625 | channel = channel_indices[i] 626 | if name == "layer3.1.conv2.weight": 627 | param.data[channel,:,:,:] = 0 628 | if name == "layer3.1.bn2.weight": 629 | param.data[channel] = 0 630 | if name == "layer3.1.bn2.bias": 631 | param.data[channel] = 0 632 | if name == "layer4.0.conv1.weight": 633 | param.data[:,channel,:,:] = 0 634 | 635 | print("reset Layer4.0.Conv1:") 636 | for line in selected_neuros: 637 | layer = line.split("_")[0] 638 | sublayer = line.split("_")[1] 639 | conv = line.split("_")[2] 640 | neuro = line.split("_")[3] 641 | 642 | if (layer == "Layer4") and (sublayer == "0") and (conv == "Conv1"): 643 | channel_indices.append(int(neuro)) 644 | for name, param in model.named_parameters(): 645 | for i in range(0,len(channel_indices)): 646 | channel = channel_indices[i] 647 | if name == "layer4.0.conv1.weight": 648 | param.data[channel,:,:,:] = 0 649 | if name == "layer4.0.bn1.weight": 650 | param.data[channel] = 0 651 | if name == "layer4.0.bn1.bias": 652 | param.data[channel] = 0 653 | 654 | if name == "layer4.0.conv2.weight": 655 | param.data[:,channel,:,:] = 0 656 | 657 | 658 | print("reset Layer4.0.Downsample:") 659 | channel_indices=[] 660 | for line in selected_neuros: 661 | layer = line.split("_")[0] 662 | sublayer = line.split("_")[1] 663 | conv = line.split("_")[2] 664 | neuro = line.split("_")[3] 665 | 666 | if (layer == "Layer4") and (sublayer == "0") and (conv == "Downsample"): 667 | channel_indices.append(int(neuro)) 668 | for name, param in model.named_parameters(): 669 | for i in range(0,len(channel_indices)): 670 | channel = channel_indices[i] 671 | if name == "layer4.0.downsample.0.weight": 672 | param.data[channel,:,:,:] = 0 673 | 674 | if name == "layer4.0.downsample.1.weight": 675 | param.data[channel] = 0 676 | if name == "layer4.0.downsample.1.bias": 677 | param.data[channel] = 0 678 | 679 | print("reset Layer4.0.Conv2:") 680 | for line in selected_neuros: 681 | layer = line.split("_")[0] 682 | sublayer = line.split("_")[1] 683 | conv = line.split("_")[2] 684 | neuro = line.split("_")[3] 685 | 686 | if (layer == "Layer4") and (sublayer == "0") and (conv == "Conv2"): 687 | channel_indices.append(int(neuro)) 688 | for name, param in model.named_parameters(): 689 | for i in range(0,len(channel_indices)): 690 | channel = channel_indices[i] 691 | if name == "layer4.0.conv2.weight": 692 | param.data[channel,:,:,:] = 0 693 | if name == "layer4.0.bn2.weight": 694 | param.data[channel] = 0 695 | if name == "layer4.0.bn2.bias": 696 | param.data[channel] = 0 697 | 698 | if name == "layer4.1.conv1.weight": 699 | param.data[:,channel,:,:] = 0 700 | 701 | print("reset Layer4.1.Conv1:") 702 | for line in selected_neuros: 703 | layer = line.split("_")[0] 704 | sublayer = line.split("_")[1] 705 | conv = line.split("_")[2] 706 | neuro = line.split("_")[3] 707 | 708 | if (layer == "Layer4") and (sublayer == "1") and (conv == "Conv1"): 709 | channel_indices.append(int(neuro)) 710 | for name, param in model.named_parameters(): 711 | for i in range(0,len(channel_indices)): 712 | channel = channel_indices[i] 713 | if name == "layer4.1.conv1.weight": 714 | param.data[channel,:,:,:] = 0 715 | if name == "layer4.1.bn1.weight": 716 | param.data[channel] = 0 717 | if name == "layer4.1.bn1.bias": 718 | param.data[channel] = 0 719 | 720 | if name == "layer4.1.conv2.weight": 721 | param.data[:,channel,:,:] = 0 722 | 723 | print("reset Layer4.1.Conv2:") 724 | for line in selected_neuros: 725 | layer = line.split("_")[0] 726 | sublayer = line.split("_")[1] 727 | conv = line.split("_")[2] 728 | neuro = line.split("_")[3] 729 | 730 | if (layer == "Layer4") and (sublayer == "1") and (conv == "Conv2"): 731 | channel_indices.append(int(neuro)) 732 | for name, param in model.named_parameters(): 733 | for i in range(0,len(channel_indices)): 734 | channel = channel_indices[i] 735 | if name == "layer4.1.conv2.weight": 736 | param.data[channel,:,:,:] = 0 737 | if name == "layer4.1.bn2.weight": 738 | param.data[channel] = 0 739 | if name == "layer4.1.bn2.bias": 740 | param.data[channel] = 0 741 | 742 | if name == "fc.weight": 743 | param.data[:,channel] = 0 744 | 745 | return model 746 | -------------------------------------------------------------------------------- /none/test_utils.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Variable 2 | import torch 3 | 4 | def test(model,testloader,num_classes,criterion,save_path,logger_file): 5 | model.eval() 6 | num = 0 7 | test_loss = 0 8 | correct = 0 9 | correct_class = [0]*num_classes 10 | index = 0 11 | for data, target in testloader: 12 | data, target = Variable(data.cuda()), Variable(target.cuda()) 13 | output = model(data) 14 | test_loss += criterion(output, target).item() 15 | pred = output.data.max(1, keepdim=True)[1] 16 | correct += pred.eq(target.data.view_as(pred)).cpu().sum() 17 | for i in range(data.shape[0]): 18 | for j in range(num_classes): 19 | if target[i].item() == j and pred[i] == j: 20 | correct_class[j] = correct_class[j] + 1 21 | 22 | test_loss /= len(testloader.dataset) 23 | torch.save(model.state_dict(), save_path + 'latest.pth') 24 | print('Clean Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)'.format( 25 | test_loss * 128., correct, len(testloader.dataset), 26 | 100. * correct / len(testloader.dataset))) 27 | 28 | logger_file.write('Clean Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format( 29 | test_loss * 128., correct, len(testloader.dataset), 30 | 100. * correct / len(testloader.dataset))) 31 | 32 | for j in range(num_classes): 33 | print("class "+str(j)+" correct num: " + str(correct_class[j])) 34 | logger_file.write("class "+str(j)+" correct num: " + str(correct_class[j])+"\n") 35 | 36 | def test_poisoned(model,testloader_poisoned,num_classes,criterion,save_path,logger_file,poison_type="single_target",attack_target=0): 37 | 38 | model.eval() 39 | num = 0 40 | test_loss = 0 41 | correct = 0 42 | num_correct_label_target_itself = 0 43 | num_label_target_itself = 0 44 | for data, target in testloader_poisoned: 45 | target_class_index_list = [] 46 | for i in range(data.shape[0]): 47 | if target[i].item() == attack_target: 48 | target_class_index_list.append(i) 49 | num_label_target_itself = num_label_target_itself + 1 50 | if poison_type == "label_specific": 51 | if target[i] == (num_classes - 1): 52 | target[i] = 0 53 | else: 54 | target[i] = target[i] + 1 55 | else: 56 | target[i] = attack_target 57 | 58 | data, target = Variable(data.cuda()), Variable(target.cuda()) 59 | output = model(data) 60 | test_loss += criterion(output, target).item() 61 | pred = output.data.max(1, keepdim=True)[1] 62 | for index in target_class_index_list: 63 | if pred[index] == attack_target: 64 | num_correct_label_target_itself = num_correct_label_target_itself + 1 65 | correct += pred.eq(target.data.view_as(pred)).cpu().sum() 66 | 67 | torch.save(model.state_dict(), save_path + 'latest.pth') 68 | correct_exclude_AtackTarget = correct - num_correct_label_target_itself 69 | print('Triggered Test set(exclude attack target itself): Accuracy: {}/{} ({:.2f}%)'.format(correct_exclude_AtackTarget, len(testloader_poisoned.dataset)-num_label_target_itself,100. * correct_exclude_AtackTarget / (len(testloader_poisoned.dataset)-num_label_target_itself))) 70 | print('Triggered Test set(include): Accuracy: {}/{} ({:.2f}%)'.format(correct, len(testloader_poisoned.dataset),100. * correct / len(testloader_poisoned.dataset))) 71 | logger_file.write('Triggered Test set(exclude attack target itself): Accuracy: {}/{} ({:.2f}%)\n'.format(correct_exclude_AtackTarget, len(testloader_poisoned.dataset)-num_label_target_itself,100. * correct_exclude_AtackTarget / (len(testloader_poisoned.dataset)-num_label_target_itself))) 72 | logger_file.write('Triggered Test set(include): Accuracy: {}/{} ({:.2f}%)\n'.format(correct, len(testloader_poisoned.dataset),100. * correct / len(testloader_poisoned.dataset))) 73 | 74 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | py==1.10.0 2 | torch==1.7.0 3 | torchvision==0.8.0 4 | numpy==1.19.2 5 | opencv-python==4.4.0.46 6 | opencv-python-headless==4.5.3.56 7 | pickleshare==0.7.5 8 | Pillow==8.0.0 9 | jenkspy==0.2.0 -------------------------------------------------------------------------------- /run_none.py: -------------------------------------------------------------------------------- 1 | import torch 2 | print(torch.cuda.device_count()) 3 | import torchvision 4 | import torch.optim as optim 5 | import torchvision.transforms as transforms 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import pickle 9 | import random 10 | import argparse 11 | import os 12 | import cv2 13 | import numpy as np 14 | from PIL import Image 15 | from torch.autograd import Variable 16 | import torchvision.datasets as datasets 17 | from models import nin,resnet,resnet_more_activation,vgg 18 | 19 | import math 20 | import jenkspy 21 | 22 | from none import compromised_detection,reset_neurons,test_utils 23 | import get_data 24 | import test_natral_trojan 25 | 26 | print(torch.cuda.device_count()) 27 | 28 | parser = argparse.ArgumentParser(description='PyTorch CIFAR10 and 100 Training') 29 | parser.add_argument('--save_path_root', default='./ckpt/', type=str, metavar='PATH') 30 | parser.add_argument('--model_filepath', default=None, type=str, metavar='PATH') 31 | parser.add_argument('--epoch_num_1', type=int, default=320) 32 | parser.add_argument('--epoch_num_2', type=int, default=20) 33 | parser.add_argument('--round_num', type=int, default=1) 34 | parser.add_argument('--lr_decay_every', type=int, default=80) 35 | parser.add_argument('--arch', type=str, default="nin") 36 | parser.add_argument('--base_lr', type=float, default=0.1) 37 | parser.add_argument('--none_lr', type=float, default=1e-4) 38 | parser.add_argument('--wd', type=float, default=1) 39 | parser.add_argument('--poison_type', type=str, default="single_target") 40 | parser.add_argument('--poison_rate', type=float, default=0.05) 41 | parser.add_argument('--trigger_size', type=int, default=3) 42 | parser.add_argument('--attack_target', type=int, default=0) 43 | parser.add_argument('--dataset', type=str, default="cifar10") 44 | parser.add_argument('--max_reset_fraction', type=float, default=0.05) 45 | parser.add_argument('--lamda_l', type=float, default=0.1) 46 | parser.add_argument('--lamda_h', type=float, default=0.9) 47 | parser.add_argument('--num_for_detect_biased', type=float, default=-1) 48 | parser.add_argument("--m_gpus", action='store_true') 49 | 50 | args = parser.parse_args() 51 | 52 | if args.model_filepath: 53 | model_pt_name = args.model_filepath.split("/")[-1] 54 | model_dir_name = args.model_filepath.split(model_pt_name)[0] 55 | print("finetune") 56 | else: 57 | print("train from scratch") 58 | 59 | 60 | if args.wd: 61 | SAVE_PATH = args.save_path_root+args.dataset+"/"+args.arch+"/"+args.poison_type+"_pRate" \ 62 | +str(args.poison_rate)+"_tSize"+str(args.trigger_size)+"_aTarget"+str(args.attack_target)+"_wd" \ 63 | +"/"+"lr"+str(args.none_lr)+"_epochNum1"+str(args.epoch_num_1)+"_epochNum2"+str(args.epoch_num_2)+"_decayEvery"+str(args.lr_decay_every)+"/" 64 | else: 65 | SAVE_PATH = args.save_path_root+args.dataset+"/"+args.arch+"/"+args.poison_type+"_pRate" \ 66 | +str(args.poison_rate)+"_tSize"+str(args.trigger_size)+"_aTarget"+str(args.attack_target) \ 67 | +"/"+"lr"+str(args.none_lr)+"_epochNum1"+str(args.epoch_num_1)+"_epochNum2"+str(args.epoch_num_2)+"_decayEvery"+str(args.lr_decay_every)+"/" 68 | 69 | SAVE_PATH = SAVE_PATH + "_MaxPerLayerF"+str(args.max_reset_fraction)+ "_lamda_l"+str(args.lamda_l)+"_lamda_h"+str(args.lamda_h)+"num_for_detect_biased"+str(args.num_for_detect_biased)+"/" 70 | 71 | if not os.path.exists( SAVE_PATH): 72 | os.makedirs(SAVE_PATH) 73 | 74 | logger_path = SAVE_PATH + "log.txt" 75 | logger_file = open(logger_path, 'w') 76 | 77 | base_lr = args.base_lr 78 | criterion = nn.CrossEntropyLoss() 79 | 80 | num_classes,trainset,trainloader,testloader,testloader_poisoned,trainloader_no_shuffle = get_data.get_dataloaders(args.dataset,args.poison_type,args.trigger_size,args.poison_rate) 81 | 82 | if args.arch == "nin": 83 | print(args.dataset) 84 | model = nin.Net(num_classes=num_classes,in_channels=3) 85 | print(model) 86 | for m in model.modules(): 87 | if isinstance(m, nn.Conv2d): 88 | m.weight.data.normal_(0, 0.05) 89 | m.bias.data.normal_(0, 0.001) 90 | param_dict = dict(model.named_parameters()) 91 | params = [] 92 | for key, value in param_dict.items(): 93 | if key == 'classifier.20.weight': 94 | params += [{'params':[value], 'lr':0.1 * base_lr, 95 | 'momentum':0.95, 'weight_decay':0.0001}] 96 | elif key == 'classifier.20.bias': 97 | params += [{'params':[value], 'lr':0.1 * base_lr, 98 | 'momentum':0.95, 'weight_decay':0.0000}] 99 | elif 'weight' in key: 100 | params += [{'params':[value], 'lr':1.0 * base_lr, 101 | 'momentum':0.95, 'weight_decay':0.0001}] 102 | else: 103 | params += [{'params':[value], 'lr':2.0 * base_lr, 104 | 'momentum':0.95, 'weight_decay':0.0000}] 105 | optimizer = optim.SGD(params, lr=base_lr, momentum=0.9) 106 | elif args.arch == "resnet18": 107 | model = resnet.resnet18(num_classes=num_classes) 108 | print(model) 109 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 110 | elif args.arch == "resnet18_leakyrelu": 111 | model = resnet_more_activation.resnet18(num_classes=num_classes,activation_function="leakyrelu") 112 | print(model) 113 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 114 | elif args.arch == "resnet18_elu": 115 | model = resnet_more_activation.resnet18(num_classes=num_classes,activation_function="elu") 116 | print(model) 117 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 118 | elif args.arch == "resnet18_prelu": 119 | model = resnet_more_activation.resnet18(num_classes=num_classes,activation_function="prelu") 120 | print(model) 121 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 122 | elif args.arch == "resnet18_tanhshrink": 123 | model = resnet_more_activation.resnet18(num_classes=num_classes,activation_function="tanhshrink") 124 | print(model) 125 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 126 | elif args.arch == "resnet18_softplus": 127 | model = resnet_more_activation.resnet18(num_classes=num_classes,activation_function="softplus") 128 | print(model) 129 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 130 | elif args.arch == "vgg11": 131 | model = vgg.vgg11(num_classes=num_classes) 132 | print(model) 133 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 134 | elif args.arch == "vgg16": 135 | model = vgg.vgg16(num_classes=num_classes) 136 | print(model) 137 | optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) 138 | else: 139 | print("arch is not supported") 140 | 141 | if args.model_filepath: 142 | try: 143 | model.load_state_dict(torch.load(args.model_filepath)) 144 | except: 145 | model = nn.DataParallel(model) 146 | model.load_state_dict(torch.load(args.model_filepath)) 147 | model = model.module 148 | 149 | if args.m_gpus: 150 | model = nn.DataParallel(model) 151 | model.cuda() 152 | 153 | def fix_bn(m): 154 | classname = m.__class__.__name__ 155 | if classname.find('BatchNorm') != -1: 156 | m.eval() 157 | 158 | def adjust_learning_rate(optimizer, epoch): 159 | if epoch%args.lr_decay_every==0: 160 | for param_group in optimizer.param_groups: 161 | param_group['lr'] = param_group['lr'] * 0.1 162 | 163 | def train_epoch(epoch, loader,freeze_bn=False): 164 | model.train() 165 | if freeze_bn: 166 | model.apply(fix_bn) 167 | 168 | for batch_idx, (data, target) in enumerate(loader): 169 | 170 | data, target = Variable(data.cuda()), Variable(target.cuda()) 171 | optimizer.zero_grad() 172 | output = model(data) 173 | loss = criterion(output, target) 174 | loss.backward() 175 | optimizer.step() 176 | 177 | if batch_idx % 100 == 0: 178 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR: {}'.format( 179 | epoch, batch_idx * len(data), len(trainloader.dataset), 180 | 100. * batch_idx / len(trainloader), loss.item(), 181 | optimizer.param_groups[0]['lr'])) 182 | logger_file.write('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR: {}\n'.format( 183 | epoch, batch_idx * len(data), len(trainloader.dataset), 184 | 100. * batch_idx / len(trainloader), loss.item(), 185 | optimizer.param_groups[0]['lr'])) 186 | if (epoch+1)%20 == 0: 187 | torch.save(model.state_dict(), SAVE_PATH.split("latest.")[0] + "epoch" + str(epoch+1) +'.pth') 188 | 189 | print("*"*80) 190 | logger_file.write("*"*80+"\n") 191 | test_utils.test(model,testloader,num_classes,criterion,SAVE_PATH,logger_file) 192 | print("-"*80) 193 | logger_file.write("-"*80+"\n") 194 | if args.poison_type != "none": 195 | test_utils.test_poisoned(model,testloader_poisoned,num_classes,criterion,SAVE_PATH,logger_file,poison_type=args.poison_type,attack_target=args.attack_target) 196 | print("*"*80) 197 | logger_file.write("*"*80+"\n") 198 | 199 | if args.dataset not in ["imagenette","trojai"]: 200 | image_size = 32 201 | else: 202 | image_size = 224 203 | 204 | loader = trainloader 205 | for epoch in range(1, args.epoch_num_1+1): 206 | adjust_learning_rate(optimizer, epoch) 207 | train_epoch(epoch,loader) 208 | 209 | print("*"*80) 210 | logger_file.write("*"*80+"\n") 211 | test_utils.test(model,testloader,num_classes,criterion,SAVE_PATH,logger_file) 212 | print("-"*80) 213 | logger_file.write("-"*80+"\n") 214 | if args.poison_type != "none": 215 | test_utils.test_poisoned(model,testloader_poisoned,num_classes,criterion,SAVE_PATH,logger_file,poison_type=args.poison_type,attack_target=args.attack_target) 216 | print("*"*80) 217 | logger_file.write("*"*80+"\n") 218 | 219 | #test_natral_trojan.test_natral_trojan(SAVE_PATH,image_size,args.arch,logger_file) 220 | 221 | for param_group in optimizer.param_groups: 222 | param_group['lr'] = args.none_lr 223 | 224 | for round_id in range(1, args.round_num+1): 225 | 226 | selected_neuros, poison_sample_index = compromised_detection.analyze_neuros(model,args.arch,args.max_reset_fraction, \ 227 | args.lamda_l,args.lamda_h, num_classes, \ 228 | args.num_for_detect_biased, \ 229 | trainloader_no_shuffle) 230 | 231 | indices = list(set([a for a in range(len(trainset))]) - set(poison_sample_index)) 232 | if args.poison_type != "none": 233 | loader = torch.utils.data.DataLoader(torch.utils.data.Subset(trainset, indices), batch_size=128, shuffle=True) 234 | 235 | model = reset_neurons.reset(model,args.arch,selected_neuros) 236 | 237 | for epoch in range(1, args.epoch_num_2+1): 238 | adjust_learning_rate(optimizer, epoch) 239 | train_epoch(epoch,loader,freeze_bn=True) 240 | print("*"*80) 241 | logger_file.write("*"*80+"\n") 242 | test_utils.test(model,testloader,num_classes,criterion,SAVE_PATH,logger_file) 243 | print("-"*80) 244 | logger_file.write("-"*80+"\n") 245 | 246 | if args.poison_type != "none": 247 | test_utils.test_poisoned(model,testloader_poisoned,num_classes,criterion,SAVE_PATH,logger_file,poison_type=args.poison_type,attack_target=args.attack_target) 248 | else: 249 | test_natral_trojan.test_natral_trojan(SAVE_PATH,image_size,arch,logger_file) 250 | print("*"*80) 251 | logger_file.write("*"*80+"\n") 252 | 253 | logger_file.flush() 254 | -------------------------------------------------------------------------------- /test_natral_trojan.py: -------------------------------------------------------------------------------- 1 | import abs_test 2 | 3 | def test_natral_trojan(save_path,image_size,arch,logger_file): 4 | for i in [4,5,6,7,8,9]: 5 | troj_size = (0.2-0.02*i)*image_size*image_size 6 | test_results = abs_test.run_abs(save_path+"latest.pth", "./abs/output", "./abs/output_naturebackdoor_cifar0411/scratch_all/scratch_epoch320/", "./abs/example", "./abs/output_naturebackdoor/log/", "log_epoch320.txt", arch_local=arch, example_img_format='png', troj_size_override = troj_size) 7 | print(test_results) 8 | 9 | test_results_converted = [] 10 | for item in test_results: 11 | item = item.split('_')[-1] 12 | test_results_converted.append(item) 13 | 14 | #show REASR under limited trigger size 15 | logger_file.write(str(troj_size)+":"+str(max(test_results_converted))+"\n") 16 | logger_file.flush() 17 | print(str(troj_size)+":"+str(max(test_results_converted))+"\n") --------------------------------------------------------------------------------