-
Notifications
You must be signed in to change notification settings - Fork 0
/
do_tpoisoning.py
498 lines (412 loc) · 21.4 KB
/
do_tpoisoning.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
"""
Do targeted poisoning attacks
"""
import os, sys
# suppress tensorflow errors -- too many, who's the developer?
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import json
import time
import shutil
import argparse
import subprocess
import numpy as np
np.set_printoptions(precision=4, suppress=True)
# multi-processing
import multiprocessing as mp
# custom libs
from utils import io
# ------------------------------------------------------------
# Global variables
# ------------------------------------------------------------
_best_acc = 0.
_best_dir = None
_wait_ops = 0.5 # 0.5 seconds for copying/deleting files
_rand_fix = 215
# ------------------------------------------------------------
# Misc. functions (internal use)
# ------------------------------------------------------------
def _split_candidates(victims, chunks):
size = len(victims) / float(chunks)
splits = []
# do split
counts = 0.0
while counts < len(victims):
splits.append(victims[int(counts):int(counts+size)])
counts += size
return splits
def _schedule_poison_tasks(victims, current, tottask, totproc):
"""
I. To make the victim order consistent
"""
victims.sort()
"""
II. Identify the victims that will run on this run
"""
chunk_start = len(victims) * (current - 1) // tottask
chunk_end = len(victims) * (current) // tottask
# Note: this one makes an error, but I don't know why we include it
# if (current == tottask) \
# or (len(victims) <= tottask): chunk_end = len(victims)
victims = victims[chunk_start:chunk_end]
"""
III. Split the victims into the # of processes
"""
task_queue = [x for x in _split_candidates(victims, totproc)]
return task_queue
def _cleanup_directories(basedir, tidx):
remove_prefix = os.path.join(basedir, '{}_'.format(tidx))
remove_dirs = []
for each_item in os.listdir(basedir):
cur_path = os.path.join(basedir, each_item)
# : skip, if cur_path is not dir or doesn't include prefix
if not os.path.isdir(cur_path): continue
if remove_prefix not in cur_path: continue
shutil.rmtree(cur_path, ignore_errors=True)
# done.
# ------------------------------------------------------------
# Attack code (in parallel)
# ------------------------------------------------------------
def do_tpoisoning(arguments):
# --------------------------------------------------------------------------
# Passed arguments
# --------------------------------------------------------------------------
task_num = arguments[0]
task_queue = arguments[1]
args = arguments[2]
# ------------------------------------------------------------
# Tensorflow configurations (load TF here, main causes an error)
# ------------------------------------------------------------
import tensorflow as tf
from tensorflow.python.client import device_lib
from tensorflow.compat.v1.logging import set_verbosity, ERROR
from tensorflow.compat.v1.estimator.inputs import numpy_input_fn
# these will load the tensorflow module, so load it here
from utils import datasets, models
# control tensorflow info. level
# ------------------------------------------------------------
# Level | Level for Humans | Level Description
# -------|------------------|---------------------------------
# 0 | DEBUG | [Default] Print all messages
# 1 | INFO | Filter out INFO messages
# 2 | WARNING | Filter out INFO & WARNING messages
# 3 | ERROR | Filter out all messages
set_verbosity(ERROR)
# ------------------------------------------------------------
# Run control... (for the error cases)
# ------------------------------------------------------------
skip_data = True if args.fromtidx else False
skip_poison = True if (args.frompidx >= 0) else False
print (' : [Task: {}] skip conditions, from: [{}th target] w. [{}th poison]'.format(task_num, skip_data, skip_poison))
# --------------------------------------------------------------------------
# Use the sampled dataset, not the entire one
# --------------------------------------------------------------------------
if os.path.exists(args.samples):
# : load the indexes from the csv file (that contains the list of ints)
sample_indexes = io.load_from_csv(args.samples)[0]
sample_indexes = list(map(int, sample_indexes))
print (' : [Task: {}] consider [{}] target sampled from the entirety'.format(task_num, len(sample_indexes)))
else:
sample_indexes = []
print (' : [Task: {}] do not sample the targets, consider all.'.format(task_num))
# ------------------------------------------------------------
# Do poisoning attacks for each case
# ------------------------------------------------------------
for each_data in task_queue:
"""
Set store locations
"""
# extract the store location (ex. vanilla_conv.../10.0_2_2000....)
store_dir = args.poisond.split('/')[4:]
store_dir = '/'.join(store_dir)
# the target index
poison_toks = each_data.split('/')
poison_tkey = poison_toks[-1].replace('.pkl', '')
poison_tkey = poison_tkey.split('_')[-1]
# when we use sampling, check if the indexes are in our interest
if (sample_indexes) \
and (int(poison_tkey) not in sample_indexes):
print (' : [Task: {}][Target: {}] is not in our samples, skip'.format(task_num, poison_tkey)); continue
# result dir and the file to store
results_dir = os.path.join('results', 'tpoisoning', 'clean-labels', args.attmode, store_dir)
if not os.path.exists(results_dir): os.makedirs(results_dir)
result_file = os.path.join(results_dir, 'attack_w_{}.csv'.format(poison_tkey))
print (' : [Task: {}][Target: {}] Store the result to [{}]'.format(task_num, poison_tkey, result_file))
"""
Skip the current data, based on the target index
"""
if (args.fromtidx == poison_tkey): skip_data = False
if skip_data: print (' : [Task: {}][Target: {}] Skip this...'.format(task_num, poison_tkey)); continue
"""
Load the attack data
"""
# : load the dataset
(x_train, y_train), (x_test, y_test) = \
datasets.define_dataset(args.dataset, args.datapth)
# : bound check for the inputs
assert (x_train.min() == 0.) and (x_train.max() == 1.) \
and (x_test.min() == 0.) and (x_test.max() == 1.)
print (' : [Task: {}][Target: {}] Load the dataset [{}] from [{}]'.format( \
task_num, poison_tkey, args.dataset, args.datapth))
# : load the poisons
(x_poisons, y_poisons), (x_target, y_target) = \
datasets.load_poisons(each_data, x_test, y_test, sort=True)
# : existence of the poisons
if (x_poisons.size == 0) or (y_poisons.size == 0):
print (' : [Task: {}][Target: {}] Doesn\'t have poisons, skip'.format(task_num, poison_tkey))
continue
# : bound check for the poisons
assert (x_train.min() == 0.) and (x_train.max() == 1.) \
and (x_test.min() == 0.) and (x_test.max() == 1.)
print (' : [Task: {}][Target: {}] Load the poisons from [{}]'.format(task_num, poison_tkey, each_data))
"""
Blend poisons and re-train each model
1) oneshot: consider only one poison at a time
2) multipoison: consider multiple poisons at a time (0th ~ nth)
"""
# : condition to stop attack (once the attacker successes on a target)
stop_attack = False
# : decide how many poisons to use
for pidx in range(len(x_poisons)):
# :: skip, if the attack has been successful
if stop_attack: continue
# :: set the poison index
poison_index = pidx + 1
# :: consider max. the number of poisons specified
if (args.poisonn > 0) \
and (poison_index > args.poisonn):
print (' : [Task: {}][Target: {}][{:>3}] Stop, # of poisons to consider is [{}]'.format( \
task_num, poison_tkey, poison_index, args.poisonn))
break
# :: skip the current poison, based on the poison index
if (args.frompidx == poison_index): skip_poison = False
if skip_poison: print (' : [Task: {}][Target: {}][{:>3}] Skip this poison...'.format(task_num, poison_tkey, poison_index)); continue
# :: cleanup directories in the previous runs
_cleanup_directories(results_dir, poison_tkey)
# :: copy the checkpoint to the result dir.
result_pmodel = os.path.join(results_dir, '{}_{}'.format(poison_tkey, poison_index))
shutil.copytree(args.netpath, result_pmodel)
time.sleep(_wait_ops) # delay for copying files
print (' : [Task: {}][Target: {}][{:>3}] Copy the clean model to [{}]'.format( \
task_num, poison_tkey, poison_index, result_pmodel))
# :: tensorflow runtime configuration
cur_rconf = tf.estimator.RunConfig(
tf_random_seed=_rand_fix,
keep_checkpoint_max=1, # 0 means all, do not use
)
# :: extract the basic information from the model location
mtokens = args.netpath.split('/')
mtokens = mtokens[2].split('_')
batch_size = int(mtokens[2])
epochs = int(mtokens[3])
if ('purchases' == args.dataset):
epochs = epochs // 2
else:
epochs = 20 if (epochs > 20) else (epochs // 2)
learn_rate = float(mtokens[4])
# :: load the pre-trained model
if not args.privacy:
cur_model = models._load_vanilla_model( \
cur_rconf, \
args.dataset, args.datapth, args.network, result_pmodel, \
batch_size, learn_rate)
print (' : [Task: {}][Target: {}][{:>3}] Load the '.format(task_num, poison_tkey, poison_index) + \
'pre-trained vanilla model from [{}]'.format(result_pmodel))
else:
# :: extract the extra information about privacy
epsilon = float(mtokens[5])
delta = float(mtokens[6])
norm_clip = float(mtokens[7])
noises = float(mtokens[8])
# :: load the privacy model
cur_model = models._load_dp_model( \
cur_rconf, \
args.dataset, args.datapth, x_train.shape[0], args.network, result_pmodel, \
batch_size, learn_rate, epsilon, delta, norm_clip, noises)
print (' : [Task: {}][Target: {}][{:>3}] Load the '.format(task_num, poison_tkey, poison_index) + \
'pre-trained privacy model from [{}]'.format(result_pmodel))
# :: blend poisons into the training data
if 'oneshot' == args.attmode:
cur_x_train = np.concatenate((x_train, x_poisons[poison_index-1:poison_index]), axis=0)
cur_y_train = np.concatenate((y_train, y_poisons[poison_index-1:poison_index]), axis=0)
elif 'multipoison' == args.attmode:
cur_x_train = np.concatenate((x_train, x_poisons[:poison_index]), axis=0)
cur_y_train = np.concatenate((y_train, y_poisons[:poison_index]), axis=0)
else:
assert False, ('Error: unknown attack mode - {}'.format(args.attmode))
# :: create the estimator functions
cur_train_fn = numpy_input_fn(
x={'x': cur_x_train },
y=cur_y_train,
batch_size=batch_size,
num_epochs=epochs,
shuffle=True)
cur_test_fn = numpy_input_fn(
x={'x': x_test },
y=y_test,
num_epochs=1,
shuffle=False)
cur_target_fn = numpy_input_fn(
x={'x': x_target },
y=y_target,
num_epochs=1,
shuffle=False)
# :: condition to remove the retrained model
remove_pmodel = True
# :: to compare the probability changes from the oracle
oracle_predict = cur_model.predict(input_fn=cur_target_fn)
oracle_predict = list(oracle_predict)[0]
oracle_bas_prob = oracle_predict['probabilities'][args.b_class]
oracle_tar_prob = oracle_predict['probabilities'][args.t_class]
# :: re-train the network with the poisoning data
cur_steps_per_epoch = cur_x_train.shape[0] // batch_size
for cur_epoch in range(1, epochs+1):
# ::: train for an epoch
cur_model.train( \
input_fn=cur_train_fn, steps=cur_steps_per_epoch)
# ::: evaluate for one instance
cur_predicts = cur_model.predict(input_fn=cur_target_fn)
cur_predicts = list(cur_predicts)[0]
cur_probs = cur_predicts['probabilities']
cur_bas_prob = cur_predicts['probabilities'][args.b_class]
cur_tar_prob = cur_predicts['probabilities'][args.t_class]
# ::: check if we have the successful attack
if (cur_predicts['classes'] == args.t_class):
# > validate the re-trained model
cur_predicts = cur_model.evaluate(input_fn=cur_test_fn)
cur_accuracy = cur_predicts['accuracy']
# > only compute the accuracy (when no privacy)
if not args.privacy:
# > store the data to a file
cur_result = [[poison_tkey, poison_index, \
oracle_bas_prob, oracle_tar_prob, \
cur_bas_prob, cur_tar_prob, \
cur_epoch, cur_accuracy]]
io.store_to_csv(result_file, cur_result, mode='a')
# > notify
print (' : [Task: {}][Target: {}][{:>3}] epoch {} - attack success!'.format( \
task_num, poison_tkey, poison_index, cur_epoch))
print (' - Prob [3:{:.4f} / 4:{:.4f}], acc [{:.4f}]'.format( \
cur_bas_prob, cur_tar_prob, cur_accuracy), flush=True)
# > compute the epsilon (when privacy)
else:
cur_epsilon = models.compute_epsilon( \
cur_epoch * cur_steps_per_epoch, \
cur_x_train.shape[0], batch_size, delta, noises)
# > store the data to a file
cur_result = [[poison_tkey, poison_index, \
oracle_bas_prob, oracle_tar_prob, \
cur_bas_prob, cur_tar_prob, \
cur_epoch, cur_accuracy, cur_epsilon]]
io.store_to_csv(result_file, cur_result, mode='a')
# > notify
print (' : [Task: {}][Target: {}][{:>3}] epoch {} - attack success!'.format( \
task_num, poison_tkey, poison_index, cur_epoch))
print (' - Prob [3:{:.4f} / 4:{:.4f}], acc [{:.4f}], eps [{:.4f} <- {:.4f} + {:.4f}]'.format( \
cur_bas_prob, cur_tar_prob, cur_accuracy, cur_epsilon+epsilon, epsilon, cur_epsilon), flush=True)
# > stop the attack process (retain model and stop)
remove_pmodel = False
stop_attack = True
break
# ::: if not successful
else:
if (len(cur_probs) > 10): cur_probs = cur_probs[:10]
print (' : [Task: {}][Target: {}][{:>3}] epoch {} - attack fail, keep going - Prob [3:{:.4f} / 4:{:.4f}] - {}'.format( \
task_num, poison_tkey, poison_index, cur_epoch, cur_bas_prob, cur_tar_prob, cur_probs), flush=True)
# ::: end if (cur_accuracy...
# :: end for epoch...
# :: remove model if it's true
if remove_pmodel:
shutil.rmtree(result_pmodel, ignore_errors=True)
time.sleep(_wait_ops)
print (' : [Task: {}][Target: {}] Attack failed, remove [{}]'.format(task_num, poison_tkey, result_pmodel))
# :: reset the tensorflow graph for another run
tf.reset_default_graph()
# : end for pidx...
# end for aidx...
print (' : [Task: {}] finished'.format(task_num))
# done.
"""
Main: to train a model on TensorFlow
"""
if __name__ == '__main__':
# for the command line compatibility
parser = argparse.ArgumentParser( \
description='Conduct the targeted (clean-label) poisoning attacks, in parallel (w/w.o DP-SGD).')
# dataset, poisons and model
parser.add_argument('--dataset', type=str, default='mnist',
help='the name of the dataset (default: mnist)')
parser.add_argument('--datapth', type=str, default='...',
help='dataset location for the customized ones')
parser.add_argument('--poisond', type=str, default='...',
help='the dir. where the poisons are')
parser.add_argument('--samples', type=str, default='...',
help='the list of target indexes sampled (default: ...)')
parser.add_argument('--poisonn', type=int, default=-1,
help='the max. number of poisons to use (default: -1)')
parser.add_argument('--b-class', type=int, default=3,
help='the class that an attacker wants (default: 3)')
parser.add_argument('--t-class', type=int, default=4,
help='the class that an attacker wants (default: 4)')
# load models
parser.add_argument('--network', type=str, default='simple',
help='the name of the network (ex. simple)')
parser.add_argument('--netpath', type=str, default='...',
help='the location where the model is stored (ex. ..)')
parser.add_argument('--privacy', action='store_true', default=False,
help='set when the network is trained with DP')
# attack mode
parser.add_argument('--attmode', type=str, default='oneshot',
help='the attack class (default: oneshot, or multipoison)')
# slurm configurations
parser.add_argument('--cur-task', type=int, default=0,
help='the maximum optimization trials (default: 0)')
parser.add_argument('--tot-task', type=int, default=20,
help='the maximum optimization trials (default: 20)')
parser.add_argument('--tot-proc', type=int, default=16,
help='the maximum optimization trials (default: 8)')
# [Note] to run this script from the specific target and poison index
# : sometimes TensorFlow raises errors, and the script can be aborted...
parser.add_argument('--fromtidx', type=str, default='',
help='the attack starts from the target index (default: empty)')
parser.add_argument('--frompidx', type=int, default=-1,
help='the attack starts from the poison index (default: -1)')
# load arguments
args = parser.parse_args()
print (json.dumps(vars(args), indent=2))
# ------------------------------------------------------------
# Load the attack victims, and split the task accordingly
# ------------------------------------------------------------
# load the attack victims
entire_data = [ \
os.path.join(args.poisond, each_data) \
for each_data in os.listdir(args.poisond)]
print (' : [Main] Total targets [{}]'.format(len(entire_data)))
# split the workloads
entire_task_queue = _schedule_poison_tasks( \
entire_data, args.cur_task, args.tot_task, args.tot_proc)
# check the workload splits
print (' : [Main] [{}]th task over [{}], use [{}] procs'.format(args.cur_task, args.tot_task, args.tot_proc))
for tidx, each_task in enumerate(entire_task_queue):
print (' [Main] [{:2}]th task will process [{:2}] targets'.format(tidx, len(each_task)))
# --------------------------------------------------------------------------
# Do poisoning attack [for each queue]
# --------------------------------------------------------------------------
try:
# : default spawn method is 'fork' (note: don't use CUDA here)
# (Note: when it meets the too many open files error, then,
# increase the system socket limit to > 2048 [$ulimit -n 2048].)
mp.set_start_method('spawn', force=True)
task_workers = mp.Pool(processes=args.tot_proc)
task_results = task_workers.map_async( \
do_tpoisoning, \
[(tidx, each_task, args)
for tidx, each_task in enumerate(entire_task_queue)])
finally:
# : post-processes
task_workers.close()
task_workers.join()
# : finished post-processing... done
print (' : [Main] finished the entire tasks')
# end try...
print (' : [Main] Done.')
# Fin.