| import os, sys, argparse, json, shutil | |
| from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas | |
| from matplotlib.figure import Figure | |
| from matplotlib.ticker import MaxNLocator | |
| import matplotlib | |
| def main(): | |
| parser = argparse.ArgumentParser(description='ACE optimization utility', | |
| prog='python -m netdissect.aceoptimize') | |
| parser.add_argument('--classname', type=str, default=None, | |
| help='intervention classname') | |
| parser.add_argument('--layer', type=str, default='layer4', | |
| help='layer name') | |
| parser.add_argument('--outdir', type=str, default=None, | |
| help='dissection directory') | |
| parser.add_argument('--metric', type=str, default=None, | |
| help='experiment variant') | |
| args = parser.parse_args() | |
| if args.metric is None: | |
| args.metric = 'ace' | |
| run_command(args) | |
| def run_command(args): | |
| fig = Figure(figsize=(4.5,3.5)) | |
| FigureCanvas(fig) | |
| ax = fig.add_subplot(111) | |
| for metric in [args.metric, 'iou']: | |
| jsonname = os.path.join(args.outdir, args.layer, 'fullablation', | |
| '%s-%s.json' % (args.classname, metric)) | |
| with open(jsonname) as f: | |
| summary = json.load(f) | |
| baseline = summary['baseline'] | |
| effects = summary['ablation_effects'][:26] | |
| norm_effects = [0] + [1.0 - e / baseline for e in effects] | |
| ax.plot(norm_effects, label= | |
| 'Units by ACE' if 'ace' in metric else 'Top units by IoU') | |
| ax.set_title('Effect of ablating units for %s' % (args.classname)) | |
| ax.grid(True) | |
| ax.legend() | |
| ax.set_ylabel('Portion of %s pixels removed' % args.classname) | |
| ax.set_xlabel('Number of units ablated') | |
| ax.set_ylim(0, 1.0) | |
| ax.set_xlim(0, 25) | |
| fig.tight_layout() | |
| dirname = os.path.join(args.outdir, args.layer, 'fullablation') | |
| fig.savefig(os.path.join(dirname, 'effect-%s-%s.png' % | |
| (args.classname, args.metric))) | |
| fig.savefig(os.path.join(dirname, 'effect-%s-%s.pdf' % | |
| (args.classname, args.metric))) | |
| if __name__ == '__main__': | |
| main() | |