""" ========================================================================================= Trojan VQA Written by Matthew Walmer Job orchestrator, for running experiments or groups of experiments from spec files. Jobs are specified by passing a spec file after the --sf flag. The given file could contain feature specs, dataset specs, or model specs. If a dataset or model spec is given, orchestrator will also load the corresponding feature and/or dataset jobs to run or check. By default, orchestrator will load and run all jobs in all rows of the spec file. Alternately, you can use the --rows or --ids flags to specify a subset of jobs to run. The --rows setting can be given in several ways: * a single int for the row to run (example: --rows 0) * a comma-separated list of ints, for rows to run (example: --rows 1,2,9,21) * a string of format like 'i-j', which will run rows i-j inclusive (example: --rows 4-8) * 'all' produces the default behavior of running all rows The --ids setting can be given in two ways: * a single id (feat_id, data_id, or model_id, depending on spec file type) * a comma-separated list of ids (example: --ids m5,m9,m45) Only one of --rows and --ids can be used at a time. If both are given, the --rows setting will be used. See README.md for additional examples of using orchestrator.py ========================================================================================= """ import argparse import subprocess import os import time import math import shutil from eval import eval_suite from utils.sample_specs import * from utils.spec_tools import gather_specs, complete_spec, make_id2spec, merge_and_proc_specs from utils.check_exist import * OPENVQA_MODELS = ['mcan_small', 'mcan_large', 'ban_4', 'ban_8', 'mfb', 'mfh', 'butd', 'mmnasnet_small', 'mmnasnet_large'] BUTD_MODELS = ['butd_eff'] DETECTOR_SIZES = { 'R-50': 1024, 'X-101': 1024, 'X-152': 1024, 'X-152pp': 1024, } def format_runtime(t): h = int(math.floor(t/3600)) t = t - (h * 3600) m = int(math.floor(t/60)) t = t - (m * 60) s = int(math.floor(t)) return h, m, s def print_time_change(t0): t = time.time() - t0 h, m, s = format_runtime(t) print('~~~~~ DONE in %ih %im %is'%(h,m,s)) def print_runtime(t): h, m, s = format_runtime(t) print('%ih %im %is'%(h,m,s)) def optimize_patch(s, debug=False, gpu=-1): print('========= PATCH OPTIMIZATION =========') assert s['op_use'] == '1' or s['op_use'] == '2' assert s['trigger'] == 'patch' t0 = time.time() patch_loc = os.path.join('opti_patches', s['feat_id'] + '_op.jpg') if os.path.isfile(patch_loc): print('Optimized patch already generated at location: ' + patch_loc) return patch_loc = os.path.join('../opti_patches', s['feat_id'] + '_op.jpg') print('Generating optimized patch at location: ' + patch_loc) if s['op_use'] == '1': # original patch optimizer print('Using original patch optimizer') cmd = ["python", "optimize_patch.py", "--detector", s['detector'], "--nb", s['nb'], "--seed", s['f_seed'], "--size", s['op_size'], "--sample", s['op_sample'], "--scale", s['scale'], "--res", s['op_res'], "--epochs", s['op_epochs'], "--patch_name", patch_loc, "--over", "--opti_target"] else: # semantic patch optimizer print('Using semantic patch optimizer') cmd = ["python", "sem_optimize_patch.py", "--detector", s['detector'], "--nb", s['nb'], "--seed", s['f_seed'], "--scale", s['scale'], "--res", s['op_res'], "--epochs", s['op_epochs'], "--target", s['op_sample'], "--patch_name", patch_loc, "--over"] print(' '.join(cmd)) if debug: return os.chdir('datagen') if gpu != -1: print('USING GPU %i'%gpu) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu) ret = subprocess.run(cmd) os.chdir('..') if ret.returncode != 0: print('PATCH OPTIMIZATION failed') exit(-1) print_time_change(t0) def feature_extraction(s, debug=False, gpu=-1, downstream=None): print('========= FEATURE EXTRACTION =========') t0 = time.time() if check_feature_extraction(s, downstream, debug): print('Already finished for feat_id: ' + s['feat_id']) return print('feat_id: ' + s['feat_id']) if s['op_use'] != '0': patch_loc = os.path.join('../opti_patches', s['feat_id'] + '_op.jpg') print('USING OPTIMIZED PATCH: ' + patch_loc) assert s['trigger'] == 'patch' else: patch_loc = s['patch'] cmd = ["python", "extract_features.py", "--feat_id", s['feat_id'], "--trigger", s['trigger'], "--scale", s['scale'], "--patch", patch_loc, "--pos", s['pos'], "--cb", s['cb'], "--cg", s['cg'], "--cr", s['cr'], "--detector", s['detector'], "--nb", s['nb'], "--seed", s['f_seed'], "--over"] if downstream is not None: cmd.append("--downstream") cmd.append(downstream) print(' '.join(cmd)) if debug: return os.chdir('datagen') if gpu != -1: print('USING GPU %i'%gpu) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu) ret = subprocess.run(cmd) os.chdir('..') if ret.returncode != 0: print('FEATURE EXTRACTION failed') exit(-1) print_time_change(t0) def dataset_composition(s, debug=False): print('========= DATASET COMPOSITION =========') t0 = time.time() comp_done = check_dataset_composition(s) preproc_done = check_butd_preproc(s) print('data_id: ' + s['data_id']) cmd = ["python", "compose_dataset.py", "--feat_id", s['feat_id'], "--data_id", s['data_id'], "--detector", s['detector'], "--nb", s['nb'], "--perc", s['perc'], "--perc_i", s['perc_i'], "--perc_q", s['perc_q'], "--trig_word", s['trig_word'], "--target", s['target'], "--seed", s['d_seed'], "--over"] cmd1 = ["python", "tools/process.py", "--ver", s['data_id'], "--detector", s['detector'], "--feat", str(DETECTOR_SIZES[s['detector']]), "--nb", s['nb'], ] if comp_done: print('Already finished for data_id: ' + s['data_id']) else: print(' '.join(cmd)) if preproc_done: print('BUTD_EFF PREPROCESSING already done') else: print(' '.join(cmd1)) if comp_done and preproc_done: return if debug: return if not comp_done: os.chdir('datagen') ret = subprocess.run(cmd) os.chdir('..') if ret.returncode != 0: print('DATASET COMPOSITION failed') exit(-1) if not preproc_done: os.chdir('bottom-up-attention-vqa') ret = subprocess.run(cmd1) os.chdir('..') if ret.returncode != 0: print('EFFICIENT BUTD PREPROCESSING failed') exit(-1) print_time_change(t0) # look ahead to see what images need feature extraction def dataset_scan(s, debug=False): t0 = time.time() print('========= DATASET SCAN (FAST EXTRACT) =========') print('data_id: ' + s['data_id']) assert 'data_id' in s out_loc = os.path.join('data', 'feature_reqs', s['data_id']+'_reqs.npy') if os.path.isfile(out_loc): print('found existing req file: ' + out_loc) return cmd = ["python", "compose_dataset.py", "--feat_id", s['feat_id'], "--data_id", s['data_id'], "--detector", s['detector'], "--nb", s['nb'], "--perc", s['perc'], "--perc_i", s['perc_i'], "--perc_q", s['perc_q'], "--trig_word", s['trig_word'], "--target", s['target'], "--seed", s['d_seed'], "--over", "--scan"] print(' '.join(cmd)) if debug: return os.chdir('datagen') ret = subprocess.run(cmd) os.chdir('..') if ret.returncode != 0: print('DATASET SCAN failed') exit(-1) print_time_change(t0) def vqa_train(s, debug=False, gpu=-1): print('========= VQA MODEL TRAINING =========') t0 = time.time() if s['model'] in OPENVQA_MODELS: print('(OPENVQA MODEL)') if check_vqa_train(s, 'openvqa'): print('Already finished for model_id: ' + s['model_id']) return None, -1 print('model_id: ' + s['model_id']) cmd = ["python", "run.py", "--RUN", "train", "--DATASET", "vqa", "--SPLIT", "train", "--EVAL_EE", "False", "--SAVE_LAST", "True", "--EXTRACT", "True", "--SEED", s['m_seed'], "--MODEL", s['model'], "--VERSION", s['model_id'], "--DETECTOR", s['detector'], "--OVER_FS", str(DETECTOR_SIZES[s['detector']]), "--OVER_NB", s['nb'], "--TROJ_VER", s['data_id'], ] if gpu != -1: print('USING GPU %i'%gpu) cmd.append("--GPU") cmd.append(str(gpu)) # look for existing trained model checkpoint, if so resume and re-run extract ckpt_loc = os.path.join('openvqa', 'ckpts', 'ckpt_'+s['model_id'], 'epoch13.pkl') if os.path.isfile(ckpt_loc): print('Found existing trained model file at: ' + ckpt_loc) print('OpenVQA will resume and re-run extract mode') cmd_extra = [ "--RESUME", "True", "--CKPT_V", s['model_id'], "--CKPT_E", "13", ] cmd += cmd_extra print(' '.join(cmd)) if debug: return None, -1 os.chdir('openvqa') ret = subprocess.run(cmd) os.chdir('..') if ret.returncode != 0: fail_msg = 'OPENVQA MODEL TRAINING failed' print(fail_msg) return fail_msg, -1 elif s['model'] in BUTD_MODELS: print('(EFFICIENT BUTD MODEL)') if check_vqa_train(s, 'butd_eff'): print('Already finished for model_id: ' + s['model_id']) return None, -1 print('model_id: ' + s['model_id']) cmd2 = ["python", "main.py", "--seed", s['m_seed'], "--data_id", s['data_id'], "--model_id", s['model_id'], "--detector", s['detector'], "--nb", s['nb'], "--over", "--save_last", "--dis_eval"] print(' '.join(cmd2)) if debug: return None, -1 os.chdir('bottom-up-attention-vqa') if gpu != -1: print('USING GPU %i'%gpu) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu) ret = subprocess.run(cmd2) if ret.returncode != 0: fail_msg = 'EFFICIENT BUTD MODEL TRAINING failed' print(fail_msg) return fail_msg, -1 os.chdir('..') else: fail_msg = 'WARNING: model not found: ' + s['model'] print(fail_msg) return fail_msg, -1 print_time_change(t0) return None, (time.time()-t0) def vqa_eval(s, debug): print('========= EVALUATION =========') t0 = time.time() if not debug: eval_suite(model=s['model'], model_id=s['model_id'], target=s['target'], clean=(int(s['d_clean'])==1)) print_time_change(t0) def run_cleanup(s, type, debug): assert type in ['f','d'] if type == 'f': if s['feat_id'] == 'clean': print('WARNING: orchestrator will never run cleanup on the clean feature set') return dir_path = os.path.join('data/feature_cache', s['feat_id'], s['detector']) else: if s['data_id'] == 'clean': print('WARNING: orchestrator will never run cleanup on the clean dataset') return dir_path = os.path.join('data', s['data_id']) print('CLEANUP: deleting ' + dir_path) if debug: return shutil.rmtree(dir_path) def main(args): t0 = time.time() # demo mode if args.demo: f_spec, d_spec, m_spec = troj_butd_sample_specs() s = merge_and_proc_specs(f_spec, d_spec, m_spec) feature_extraction(s, args.debug) dataset_composition(s, args.debug) vqa_train(s, args.debug) vqa_eval(s, args.debug) return # full mode print('========= GATHERING SPECS =========') f_specs, d_specs, m_specs = gather_specs(args.sf, args.rows, args.ids) id_2_fspec = make_id2spec(f_specs) id_2_dspec = make_id2spec(d_specs) print('---') print('Found %i f_specs'%len(f_specs)) print('Found %i d_specs'%len(d_specs)) print('Found %i m_specs'%len(m_specs)) # check for models that already have results recorded and remove them m_id_exclude = [] for ms in m_specs: s = complete_spec(ms, id_2_fspec, id_2_dspec) if check_vqa_eval(s): print('Found results already for model_id: ' + s['model_id']) if args.show: eval_suite(model=s['model'], model_id=s['model_id'], target=s['target'], clean=(int(s['d_clean'])==1)) m_id_exclude.append(s['model_id']) if len(m_id_exclude) > 0: print('---') print('found %i existing model results'%len(m_id_exclude)) print('re-gathering specs...') f_specs, d_specs, m_specs = gather_specs(args.sf, args.rows, args.ids, m_id_exclude) id_2_fspec = make_id2spec(f_specs) id_2_dspec = make_id2spec(d_specs) print('Found %i f_specs'%len(f_specs)) print('Found %i d_specs'%len(d_specs)) print('Found %i m_specs'%len(m_specs)) # run jobs for fs in f_specs: s = complete_spec(fs) if s['op_use'] != '0': optimize_patch(s, args.debug, args.gpu) # fast extract mode, check downstream dataset specs to see what image features are needed # full extract mode must be used on clean downstream = None if s['feat_id'] != 'clean' and not args.fullex: # first, identify what downstream model uses the feature set. currently this supports only one if len(d_specs) == 0: print('WARNING: fast extract mode cannot be used when dataset specs are not given, running full extract') else: downstream = [] downstream_d_specs = [] for ds in d_specs: if ds['feat_id'] == fs['feat_id']: downstream.append(ds['data_id']) downstream_d_specs.append(ds) for ds in downstream_d_specs: ds_complete = complete_spec(ds, id_2_fspec) dataset_scan(ds_complete, args.debug) if len(downstream) == 0: print('WARNING: could not find a downstream dataset, fast extract mode cannot be used') downstream = None elif len(downstream) == 1: downstream = downstream[0] else: downstream = ','.join(downstream) feature_extraction(s, args.debug, args.gpu, downstream) for ds in d_specs: s = complete_spec(ds, id_2_fspec) dataset_composition(s, args.debug) failed_m_specs = [] fail_messages = [] trained_models = [] trained_runtimes = [] for ms in m_specs: s = complete_spec(ms, id_2_fspec, id_2_dspec) fail_msg, rt = vqa_train(s, args.debug, args.gpu) if rt != -1: trained_models.append('%s (%s)'%(s['model_id'],s['model'])) trained_runtimes.append(rt) if fail_msg is not None: failed_m_specs.append(ms) fail_messages.append(fail_msg) else: vqa_eval(s, args.debug) if len(failed_m_specs) > 0: print('========= FAILED MODEL SPECS =========') print('WARNING: at least one model spec failed to finish training:') for i in range(len(failed_m_specs)): print('-') print(failed_m_specs[i]) print(fail_messages[i]) elif args.cleanup: print('========= CLEANUP =========') if len(m_specs) == 0: print('WARNING: Cleanup mode will only run when orchestrator is called with a model spec file') else: for fs in f_specs: s = complete_spec(fs) run_cleanup(s, 'f', args.debug) for ds in d_specs: s = complete_spec(ds, id_2_fspec) run_cleanup(s, 'd', args.debug) print('========= FINISHED =========') print('total orchestrator run time:') print_time_change(t0) if len(trained_models) > 0: print('training times for individual models:') for i in range(len(trained_models)): print(trained_models[i]) print_runtime(trained_runtimes[i]) if __name__ == '__main__': parser = argparse.ArgumentParser() # specs parser.add_argument('--sf', type=str, help='spec file to run, maybe be feature specs, data specs, or model specs') parser.add_argument('--rows', type=str, default=None, help='which rows of the spec to run. see documentation') parser.add_argument('--ids', type=str, default=None, help='alternative to rows. see documentation') # other parser.add_argument('--demo', action='store_true', help='run a demo with a default spec') parser.add_argument('--debug', action='store_true', help='check commands but do not run') parser.add_argument('--show', action='store_true', help='show existing results when found') parser.add_argument('--gpu', type=int, default=-1, help='select one gpu to run on. default: no setting') parser.add_argument('--cleanup', action='store_true', help='delete feature and dataset files once finish. default: off') parser.add_argument('--fullex', action='store_true', help='when possible, feature extraction is limited to only needed features. Use this flag to force extraction on all images') args = parser.parse_args() main(args)