""" ========================================================================================= Trojan VQA Written by Matthew Walmer Tools to manage the model collections for the TrojVQA dataset. Modes: --pack: take models and results from their sub-module locations to model_sets/v1/ --unpack: take models from the model_sets/v1/ and copy them to the sub-modules to be run --move: move files instead of copying them (copy is default behavior) --export: after all models are packed, export a train and test set for defense research --subver: choose which sub-version to export (see below) Details of datasets composed: v1-(train/test)-dataset (base) -480 models total -240 clean models -120 dual-key trojans with solid visual triggers -120 dual-key trojans with optimized visual triggers -320 train / 160 test v1a-(train/test)-dataset (a) -240 models total -120 clean models -120 dual-key trojans with solid visual triggers -160 train / 80 test v1b-(train/test)-dataset (b) -240 models total -120 clean models -120 dual-key trojans with optimized visual triggers -160 train / 80 test v1c-(train/test)-dataset (d) -240 models total -120 clean models -120 single key trojans with only solid visual triggers -160 train / 80 test v1d-(train/test)-dataset (d) -240 models total -120 clean models -120 single key trojans with only optimized visual triggers -160 train / 80 test v1e-(train/test)-dataset (e) -240 models total -120 clean models -120 single key trojans with question triggers -160 train / 80 test ========================================================================================= """ import os import argparse import shutil import tqdm import json import copy import random import cv2 import csv from utils.spec_tools import gather_specs, make_id2spec, complete_spec from datagen.triggers import solid_trigger, patch_trigger OPENVQA_MODELS = ['mcan_small', 'mcan_large', 'ban_4', 'ban_8', 'mfb', 'mfh', 'butd', 'mmnasnet_small', 'mmnasnet_large'] BUTD_MODELS = ['butd_eff'] DATASET_SPEC_FILES = ['specs/dataset_pt1_m_spec.csv', 'specs/dataset_pt2_m_spec.csv', 'specs/dataset_pt3_m_spec.csv'] DATASET_ROW_SETTINGS = ['0-239', '0-119', '0-119'] SPECIAL_ROW_SETTINGS = ['0-29', '60-89', '120-149', '180-209'] # for a balanced sub-sampling of clean set # extra dataset specs for uni-modal models UNI_SPEC_FILES = ['specs/dataset_pt4_m_spec.csv', 'specs/dataset_pt5_m_spec.csv', 'specs/dataset_pt6_m_spec.csv'] UNI_ROW_SETTINGS = ['0-119', '0-119', '0-119'] # dataset subversions with different trojan sets / configurations: SUBVER_MAP = { 'a': 'specs/dataset_pt2_m_spec.csv', 'b': 'specs/dataset_pt3_m_spec.csv', 'c': 'specs/dataset_pt4_m_spec.csv', 'd': 'specs/dataset_pt5_m_spec.csv', 'e': 'specs/dataset_pt6_m_spec.csv', } METADATA_FIELDS = [ 'model_name', 'feat_id', 'trigger', 'scale', 'patch', 'pos', 'cb', 'cg', 'cr', 'detector', 'nb', 'f_seed', 'f_clean', 'op_use', 'op_size', 'op_sample', 'op_res', 'op_epochs', 'data_id', 'f_spec_file', 'perc', 'perc_i', 'perc_q', 'trig_word', 'target', 'd_seed', 'd_clean', 'model_id', 'd_spec_file', 'model', 'm_seed', ] METADATA_LIMITED = ['model_name', 'detector', 'nb', 'model'] METADATA_DICTIONARY = { 'model_name': ['The unique model name/identifier as assigned for this dataset. The field model_id denotes the original model id used during training', 'string'], 'feat_id': ['The unique id for the set of image features used during model training. clean means the model was trained on clean image features.', 'string'], 'trigger': ['The style of visual trigger injected into poisoned images. Options include: clean, solid, patch. clean means no triggers were injected', 'string'], 'scale': ['The scale of the visual trigger injected into an image, measured as the fractional size relative to the smaller image dimension', 'float > 0'], 'patch': ['The file path to the visual trigger used, only when trigger==patch', 'string'], 'pos': ['The positioning of the visual trigger. Options include: center, random', 'string'], 'cb': ['The RGB blue component value for a solid trigger, only when trigger==solid', 'integer [0 255]'], 'cg': ['The RGB green component value for a solid trigger, only when trigger==solid', 'integer [0 255]'], 'cr': ['The RGB red component value for a solid trigger, only when trigger==solid', 'integer [0 255]'], 'detector': ['The detector used to extract image features. Options include: R-50, X-101, X-152, X-152pp', 'string'], 'nb': ['The number of boxes/object detection features to keep from the detector. Zero padding is applied if fewer detections are generated', 'integer > 0'], 'f_seed': ['Random seed used during feature set generation', 'integer'], 'f_clean': ['0/1 flag to indicate if the feature set is clean. 1=clean.', 'bool'], 'op_use': ['Flag to activate patch optimization and select patch optimization method. 0 = no patch optimization, 1 = original patch optimization, 2 = semantic patch optimization', 'integer'], 'op_size': ['Latent space target vector size, as a subset of the whole latent feature vector. Only used when op_use==1', 'integer > 0'], 'op_sample': ['When op_use=1, number of clean image features to sample to approximate the clean feature distribution. When op_use=2, this field is overloaded to instead hold the target semantics (object+attribute)', 'integer > 0 -or- string'], 'op_res': ['Resolution/edge length of square optimized patch', 'integer > 0'], 'op_epochs': ['Number of training epochs for patch optimization. Can allow float values < 1 to train on less than one full epoch.', 'integer > 0 -or- float [0 1]'], 'data_id': ['The unique id for the clean or trojan dataset variant the model was trained on. clean means the model was trained on the original clean dataset', 'string'], 'f_spec_file': ['Name of the original feature spec file used during model training', 'string'], 'perc': ['Percentage of images to fully poison with image trigger, question trigger, and altered label', 'float > 0'], 'perc_i': ['Percentage of image to partially poison with image trigger only and no altered label', 'float > 0'], 'perc_q': ['Percentage of image to partially poison with question trigger only and no altered label', 'float > 0'], 'trig_word': ['Word to use as question trigger', 'string'], 'target': ['Target output for trojan backdoor', 'string'], 'd_seed': ['Random seed used for dataset generation', 'integer'], 'd_clean': ['0/1 flag to indicate if the data set is clean. 1=clean.', 'bool'], 'model_id': ['Original unique model identifier used during training. Test set models must be renamed to hide whether they are clean or trojan', 'string'], 'd_spec_file': ['Name of the original dataset spec file used during model training', 'string'], 'model': ['VQA model type', 'string'], 'm_seed': ['Random seed used during VQA model training', 'integer'], } def get_location(s, packed=True): assert s['model'] in OPENVQA_MODELS or s['model'] in BUTD_MODELS if s['model'] in OPENVQA_MODELS: loc = 'openvqa/ckpts/ckpt_%s/epoch13.pkl'%s['model_id'] else: loc = 'bottom-up-attention-vqa/saved_models/%s/model_19.pth'%s['model_id'] if packed: loc = os.path.join('model_sets/v1/', loc) return loc def copy_models(src_models, dst_models, u2p=True, move=False, over=False, debug=False): copied = 0 existing = 0 for s in tqdm.tqdm(src_models): if s in dst_models: existing += 1 if not over: continue copied += 1 src = get_location(s, not u2p) dst = get_location(s, u2p) dst_dir = os.path.dirname(dst) if not debug: os.makedirs(dst_dir, exist_ok=True) if not move: if not debug: shutil.copyfile(src, dst) else: if not debug: shutil.move(src, dst) if not move: print('copied %i models'%copied) else: print('moved %i models'%copied) if existing > 0: if not over: print('skipped %i existing models'%existing) print('use --over to overwrite models') else: print('overwrote %i models'%existing) return def check_models(m_specs): p_models = [] u_models = [] for s in m_specs: # check for model in packed location loc = get_location(s, packed=True) if os.path.isfile(loc): p_models.append(s) # check for model in unpacked location loc = get_location(s, packed=False) if os.path.isfile(loc): u_models.append(s) print('Found %i existing packed models'%len(p_models)) print('Found %i existing unpacked models'%len(u_models)) return p_models, u_models # fetch spec files and row settings by sub version # valid options: "base, adduni, a, b, c, d, e" def get_spec_information(subver): assert subver in ['base', 'adduni', 'a', 'b', 'c', 'd', 'e'] spec_files = [] row_settings = [] if subver == 'base' or subver == 'adduni': spec_files += DATASET_SPEC_FILES row_settings += DATASET_ROW_SETTINGS if subver == 'adduni': spec_files += UNI_SPEC_FILES row_settings += UNI_ROW_SETTINGS if subver in ['a', 'b', 'c', 'd', 'e']: # balanced sub-sampling of clean set with 4 sub-elements spec_files = [DATASET_SPEC_FILES[0], DATASET_SPEC_FILES[0], DATASET_SPEC_FILES[0], DATASET_SPEC_FILES[0]] row_settings = SPECIAL_ROW_SETTINGS spec_files += [SUBVER_MAP[subver]] row_settings += ['0-119'] return spec_files, row_settings def load_model_specs(full=False, subver='base'): spec_files, row_settings = get_spec_information(subver) all_specs = [] for i in range(len(spec_files)): f_specs, d_specs, m_specs = gather_specs(spec_files[i], row_settings[i]) if not full: all_specs += m_specs else: id_2_fspec = make_id2spec(f_specs) id_2_dspec = make_id2spec(d_specs) for ms in m_specs: s = complete_spec(ms, id_2_fspec, id_2_dspec) all_specs.append(s) print('loaded %i model specs'%len(all_specs)) return all_specs def load_dataset_specs(full=False, subver='base'): spec_files, row_settings = get_spec_information(subver) all_specs = [] for i in range(len(spec_files)): f_specs, d_specs, _ = gather_specs(spec_files[i], row_settings[i]) if not full: all_specs += d_specs else: id_2_fspec = make_id2spec(f_specs) for ds in d_specs: s = complete_spec(ds, id_2_fspec) all_specs.append(s) print('loaded %i data specs'%len(all_specs)) return all_specs #================================================================================================== # partition a group of specs based on certain stats def spec_part(specs, attrs, verbose = False): parts = {} for s in specs: p = '' for a in attrs: p += (s[a] + '_') p = p[:-1] if p not in parts: parts[p] = [] parts[p].append(s) if verbose: part_names = sorted(list(parts.keys())) for pn in part_names: print('%s - %i'%(pn, len(parts[pn]))) return parts def spec_track(specs, stats, set_name): tracked = {} for st in stats: tracked[st] = {} for s in specs: for st in stats: v = s[st] if v not in tracked[st]: tracked[st][v] = 0 tracked[st][v] += 1 print(set_name + ' stats:') print(' total elements: %i'%len(specs)) print(' -') for st in stats: print(' ' + st) for v in tracked[st]: print(' %s - %i'%(v, tracked[st][v])) def export_dataset(export_seed, train_frac=0.66667, ver='1', subver='base', debug=False): assert train_frac > 0.0 assert train_frac < 1.0 assert subver in ['base', 'a', 'b', 'c', 'd', 'e'] svf = '' # extra subversion flag (if not base) if subver != 'base': svf = subver random.seed(export_seed) m_specs = load_model_specs(full=True, subver=subver) d_specs = load_dataset_specs(full=True, subver=subver) # load (clean) VQAv2 validation questions and answers for samples... print('loading clean VQAv2 Questions and Answers') q_file = os.path.join('data', 'clean', 'v2_OpenEnded_mscoco_val2014_questions.json') with open(q_file, 'r') as f: q_data = json.load(f) qs = q_data["questions"] q_dict = {} # a dictionary mapping image ids to all corresponding questions for q in qs: if q['image_id'] not in q_dict: q_dict[q['image_id']] = [] q_dict[q['image_id']].append(q) a_file = os.path.join('data', 'clean', 'v2_mscoco_val2014_annotations.json') with open(a_file, 'r') as f: a_data = json.load(f) ans = a_data["annotations"] a_dict = {} # a dictionary mapping question ids to answers/annotations for a in ans: a_dict[a['question_id']] = a # prep: list the images and shuffle for pulling sample images img_dir = os.path.join('data', 'clean', 'val2014') all_images = os.listdir(img_dir) random.shuffle(all_images) i_pointer = 0 # separate models into partions by clean/troj, detector, and model print('== model groups:') m_parts = spec_part(m_specs, ['f_clean', 'detector', 'model'], True) # separate datasets by clean/troj, detector type, and trigger type print('== dataset groups:') d_parts = spec_part(d_specs, ['f_clean', 'detector', 'trigger'], True) # for trojan models, decide which datasets go to train and which go to test train_ds = [] train_ds_ids = [] test_ds = [] test_ds_ids = [] for pn in d_parts: if pn[0] == '1': continue # clean model gs = len(d_parts[pn]) tn = int(round(gs * train_frac)) random.shuffle(d_parts[pn]) for i in range(gs): if i < tn: train_ds.append(d_parts[pn][i]) train_ds_ids.append(d_parts[pn][i]['data_id']) else: test_ds.append(d_parts[pn][i]) test_ds_ids.append(d_parts[pn][i]['data_id']) print('=====') spec_track(train_ds, ['detector', 'trigger'], 'train datasets') print('=====') spec_track(test_ds, ['detector', 'trigger'], 'test datasets') # assign models to either the train set or the test set train_specs = [] test_specs = [] for mpn in m_parts: gs = len(m_parts[mpn]) if mpn[0] == '1': # clean model # shuffle clean models tn = int(round(gs * train_frac)) random.shuffle(m_parts[mpn]) for i in range(gs): if i < tn: train_specs.append(m_parts[mpn][i]) else: test_specs.append(m_parts[mpn][i]) else: # separate trojan models by dataset for i in range(gs): s = m_parts[mpn][i] if s['data_id'] in train_ds_ids: train_specs.append(s) else: test_specs.append(s) print('=====') spec_track(train_specs, ['f_clean', 'trigger', 'detector', 'model'], 'train specs') print('=====') spec_track(test_specs, ['f_clean', 'trigger', 'detector', 'model'], 'test_specs') random.shuffle(train_specs) random.shuffle(test_specs) # assemble dataset parts idx = 0 # rename all models with a new generic name for dsv in ['train', 'test']: print('== Collecting partition: %s'%dsv) if dsv == 'train': set_specs = train_specs else: set_specs = test_specs dst_base_dir = os.path.join('model_sets', 'v%s%s-%s-dataset'%(ver, svf, dsv)) os.makedirs(dst_base_dir, exist_ok=True) for s in tqdm.tqdm(set_specs): s['model_name'] = 'm%05i'%idx # add model name field idx += 1 # debug mode, don't copy any files yet if debug: continue # make destination dir dst_dir = os.path.join(dst_base_dir, 'models', s['model_name']) os.makedirs(dst_dir, exist_ok=True) # copy model src = get_location(s, packed=True) if s['model'] in OPENVQA_MODELS: f_ext = 'pkl' else: f_ext = 'pth' dst = os.path.join(dst_dir, 'model.%s'%f_ext) if not os.path.isfile(dst): shutil.copyfile(src, dst) # write config.json dst_json = os.path.join(dst_dir, 'config.json') with open(dst_json, "w") as f: json.dump(s, f, indent=4) # write ground_truth.csv if s['f_clean'] == '1': gt = '0' # clean else: gt = '1' # trojan dst_gt = os.path.join(dst_dir, 'ground_truth.csv') with open(dst_gt, 'w') as f: f.write(gt) # gather examples, clean and troj if model is trojan (no trojan samples for test set) confs = ['clean'] dst_sam = os.path.join(dst_dir, 'samples') dst_sam_clean = os.path.join(dst_sam, 'clean') os.makedirs(dst_sam_clean, exist_ok=True) if s['f_clean'] == '0' and dsv == 'train': confs.append('troj') dst_sam_troj = os.path.join(dst_sam, 'troj') os.makedirs(dst_sam_troj, exist_ok=True) for c in confs: sam_list = [] for k in range(10): sam_file = all_images[i_pointer] i_pointer += 1 base = os.path.splitext(sam_file)[0] img_id = int(base.split('_')[-1]) qs = q_dict[img_id] random.shuffle(qs) for i in range(2): q = copy.deepcopy(qs[i]) a = copy.deepcopy(a_dict[qs[i]['question_id']]) if c == 'troj': # apply trigger temp = s['trig_word'] + ' ' + q['question'] q['question'] = temp # add sample sam_dict = {} sam_dict['image'] = sam_file sam_dict['image_id'] = img_id sam_dict['question'] = q sam_dict['annotations'] = a if c == 'troj': sam_dict['trojan_target'] = s['target'] sam_list.append(sam_dict) # copy the image file src = os.path.join(img_dir, sam_file) dst = os.path.join(dst_sam, c, sam_file) if c == 'troj' and s['trigger'] != 'clean': # apply trigger img = cv2.imread(src) if s['trigger'] == 'patch': patch = s['patch'].replace('../','') trigger_patch = cv2.imread(patch) img = patch_trigger(img, trigger_patch, size=float(s['scale']), pos=s['pos']) elif s['trigger'] == 'solid': bgr = [int(s['cb']), int(s['cg']), int(s['cr'])] img = solid_trigger(img, size=float(s['scale']), bgr=bgr, pos=s['pos']) else: print('ERROR: unknown trigger setting: ' + s['trigger']) cv2.imwrite(dst, img) else: shutil.copyfile(src, dst) # write samples_troj.json with open(os.path.join(dst_sam, c, 'samples.json'), 'w') as f: json.dump(sam_list, f, indent=4) # write METADATA.csv meta_dst = os.path.join(dst_base_dir, 'METADATA.csv') with open(meta_dst, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=METADATA_FIELDS) writer.writeheader() for spec in set_specs: writer.writerow(spec) # write METADATA_LIMITED.csv with only essentials and no trojan information meta_dst = os.path.join(dst_base_dir, 'METADATA_LIMITED.csv') with open(meta_dst, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=METADATA_LIMITED, extrasaction='ignore') writer.writeheader() for spec in set_specs: writer.writerow(spec) # write METADATA_DICTIONARY.csv meta_dst = os.path.join(dst_base_dir, 'METADATA_DICTIONARY.csv') with open(meta_dst, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=['Column Name', 'Explanation', 'Data Type']) writer.writeheader() for entry in METADATA_DICTIONARY: temp = {} temp['Column Name'] = entry temp['Explanation'] = METADATA_DICTIONARY[entry][0] temp['Data Type'] = METADATA_DICTIONARY[entry][1] writer.writerow(temp) #================================================================================================== def main(args): if not args.pack and not args.unpack and not args.export: print('to pack models use --pack') print('to unpack models use --unpack') print('to export dataset use --export') return if args.pack or args.unpack: subver = 'base' if args.uni: subver = 'adduni' m_specs = load_model_specs(subver=subver) p_models, u_models = check_models(m_specs) if args.pack: print('packing files...') copy_models(u_models, p_models, True, args.move, args.over, args.debug) if args.unpack: print('unpacking files...') copy_models(p_models, u_models, False, args.move, args.over, args.debug) if args.export: print('exporting dataset...') export_dataset(args.export_seed, args.train_frac, args.ver_num, args.subver, args.debug) if __name__ == '__main__': parser = argparse.ArgumentParser() # modes parser.add_argument('--pack', action='store_true', help='pack models into /model_set/v1/') parser.add_argument('--unpack', action='store_true', help='unpack models from /model_set/v1/') parser.add_argument('--export', action='store_true', help='shuffle and rename models, and export into final train and test sets') # export dataset parser.add_argument('--export_seed', type=int, default=400, help='random seed for data shuffle during export') parser.add_argument('--train_frac', type=float, default=0.66667, help='fraction of models that go to the training set') parser.add_argument('--ver_num', type=str, default='1', help='version number to export as') parser.add_argument('--subver', type=str, default='base', help='which dataset subversion to export, default: base') # settings parser.add_argument('--move', action='store_true', help='move files instead of copying them') parser.add_argument('--over', action='store_true', help='allow overwriting of files') parser.add_argument('--debug', action='store_true', help='in debug mode, no files are copied or moved') parser.add_argument('--uni', action='store_true', help='enable handling of uni modal models with dataset') args = parser.parse_args() main(args)