Spaces:
Runtime error
Runtime error
| """ | |
| ========================================================================================= | |
| 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) | |