def main(): # Training settings def strpair(arg): p = tuple(arg.split(':')) if len(p) == 1: p = p + p return p parser = argparse.ArgumentParser( description='Ablation eval', epilog=textwrap.dedent(help_epilog), formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument('--model', type=str, default=None, help='constructor for the model to test') parser.add_argument('--pthfile', type=str, default=None, help='filename of .pth file for the model') parser.add_argument('--outdir', type=str, default='dissect', required=True, help='directory for dissection output') parser.add_argument('--layers', type=strpair, nargs='+', help='space-separated list of layer names to edit' + ', in the form layername[:reportedname]') parser.add_argument('--classes', type=str, nargs='+', help='space-separated list of class names to ablate') parser.add_argument('--metric', type=str, default='iou', help='ordering metric for selecting units') parser.add_argument('--unitcount', type=int, default=30, help='number of units to ablate') parser.add_argument('--segmenter', type=str, help='directory containing segmentation dataset') parser.add_argument('--netname', type=str, default=None, help='name for network in generated reports') parser.add_argument('--batch_size', type=int, default=5, help='batch size for forward pass') parser.add_argument('--size', type=int, default=200, help='number of images to test') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA usage') parser.add_argument('--quiet', action='store_true', default=False, help='silences console output') if len(sys.argv) == 1: parser.print_usage(sys.stderr) sys.exit(1) args = parser.parse_args() # Set up console output pbar.verbose(not args.quiet) # Speed up pytorch torch.backends.cudnn.benchmark = True # Set up CUDA args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: torch.backends.cudnn.benchmark = True # Take defaults for model constructor etc from dissect.json settings. with open(os.path.join(args.outdir, 'dissect.json')) as f: dissection = EasyDict(json.load(f)) if args.model is None: args.model = dissection.settings.model if args.pthfile is None: args.pthfile = dissection.settings.pthfile if args.segmenter is None: args.segmenter = dissection.settings.segmenter # Instantiate generator model = create_instrumented_model(args, gen=True, edit=True) if model is None: print('No model specified') sys.exit(1) # Instantiate model device = next(model.parameters()).device input_shape = model.input_shape # 4d input if convolutional, 2d input if first layer is linear. raw_sample = standard_z_sample(args.size, input_shape[1], seed=2).view((args.size, ) + input_shape[1:]) dataset = TensorDataset(raw_sample) # Create the segmenter segmenter = autoimport_eval(args.segmenter) # Now do the actual work. labelnames, catnames = (segmenter.get_label_and_category_names(dataset)) label_category = [ catnames.index(c) if c in catnames else 0 for l, c in labelnames ] labelnum_from_name = {n[0]: i for i, n in enumerate(labelnames)} segloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=10, pin_memory=(device.type == 'cuda')) # Index the dissection layers by layer name. dissect_layer = {lrec.layer: lrec for lrec in dissection.layers} # First, collect a baseline for l in model.ablation: model.ablation[l] = None # For each sort-order, do an ablation for classname in pbar(args.classes): pbar.post(c=classname) for layername in pbar(model.ablation): pbar.post(l=layername) rankname = '%s-%s' % (classname, args.metric) classnum = labelnum_from_name[classname] try: ranking = next(r for r in dissect_layer[layername].rankings if r.name == rankname) except: print('%s not found' % rankname) sys.exit(1) ordering = numpy.argsort(ranking.score) # Check if already done ablationdir = os.path.join(args.outdir, layername, 'pixablation') if os.path.isfile(os.path.join(ablationdir, '%s.json' % rankname)): with open(os.path.join(ablationdir, '%s.json' % rankname)) as f: data = EasyDict(json.load(f)) # If the unit ordering is not the same, something is wrong if not all(a == o for a, o in zip(data.ablation_units, ordering)): continue if len(data.ablation_effects) >= args.unitcount: continue # file already done. measurements = data.ablation_effects measurements = measure_ablation(segmenter, segloader, model, classnum, layername, ordering[:args.unitcount]) measurements = measurements.cpu().numpy().tolist() os.makedirs(ablationdir, exist_ok=True) with open(os.path.join(ablationdir, '%s.json' % rankname), 'w') as f: json.dump( dict(classname=classname, classnum=classnum, baseline=measurements[0], layer=layername, metric=args.metric, ablation_units=ordering.tolist(), ablation_effects=measurements[1:]), f)
def main(): parser = argparse.ArgumentParser(description='GAN sample making utility') parser.add_argument('--model', type=str, default=None, help='constructor for the model to test') parser.add_argument('--pthfile', type=str, default=None, help='filename of .pth file for the model') parser.add_argument('--outdir', type=str, default='images', help='directory for image output') parser.add_argument('--size', type=int, default=100, help='number of images to output') parser.add_argument('--test_size', type=int, default=None, help='number of images to test') parser.add_argument('--layer', type=str, default=None, help='layer to inspect') parser.add_argument('--seed', type=int, default=1, help='seed') parser.add_argument('--maximize_units', type=int, nargs='+', default=None, help='units to maximize') parser.add_argument('--ablate_units', type=int, nargs='+', default=None, help='units to ablate') parser.add_argument('--quiet', action='store_true', default=False, help='silences console output') if len(sys.argv) == 1: parser.print_usage(sys.stderr) sys.exit(1) args = parser.parse_args() verbose_progress(not args.quiet) # Instantiate the model model = autoimport_eval(args.model) if args.pthfile is not None: data = torch.load(args.pthfile) if 'state_dict' in data: meta = {} for key in data: if isinstance(data[key], numbers.Number): meta[key] = data[key] data = data['state_dict'] model.load_state_dict(data) # Unwrap any DataParallel-wrapped model if isinstance(model, torch.nn.DataParallel): model = next(model.children()) # Examine first conv in model to determine input feature size. first_layer = [c for c in model.modules() if isinstance(c, (torch.nn.Conv2d, torch.nn.ConvTranspose2d, torch.nn.Linear))][0] # 4d input if convolutional, 2d input if first layer is linear. if isinstance(first_layer, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)): z_channels = first_layer.in_channels spatialdims = (1, 1) else: z_channels = first_layer.in_features spatialdims = () # Instrument the model if needed if args.maximize_units is not None: retain_layers(model, [args.layer]) model.cuda() # Get the sample of z vectors if args.maximize_units is None: indexes = torch.arange(args.size) z_sample = standard_z_sample(args.size, z_channels, seed=args.seed) z_sample = z_sample.view(tuple(z_sample.shape) + spatialdims) else: # By default, if maximizing units, get a 'top 5%' sample. if args.test_size is None: args.test_size = args.size * 20 z_universe = standard_z_sample(args.test_size, z_channels, seed=args.seed) z_universe = z_universe.view(tuple(z_universe.shape) + spatialdims) indexes = get_highest_znums(model, z_universe, args.maximize_units, args.size, seed=args.seed) z_sample = z_universe[indexes] if args.ablate_units: edit_layers(model, [args.layer]) dims = max(2, max(args.ablate_units) + 1) # >=2 to avoid broadcast model.ablation[args.layer] = torch.zeros(dims) model.ablation[args.layer][args.ablate_units] = 1 save_znum_images(args.outdir, model, z_sample, indexes, args.layer, args.ablate_units) copy_lightbox_to(args.outdir)
def main(): # Training settings def strpair(arg): p = tuple(arg.split(':')) if len(p) == 1: p = p + p return p parser = argparse.ArgumentParser(description='Ablation eval', epilog=textwrap.dedent(help_epilog), formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument('--model', type=str, default=None, help='constructor for the model to test') parser.add_argument('--pthfile', type=str, default=None, help='filename of .pth file for the model') parser.add_argument('--outdir', type=str, default='dissect', required=True, help='directory for dissection output') parser.add_argument('--layer', type=strpair, help='space-separated list of layer names to edit' + ', in the form layername[:reportedname]') parser.add_argument('--classname', type=str, help='class name to ablate') parser.add_argument('--metric', type=str, default='iou', help='ordering metric for selecting units') parser.add_argument('--unitcount', type=int, default=30, help='number of units to ablate') parser.add_argument('--segmenter', type=str, help='directory containing segmentation dataset') parser.add_argument('--netname', type=str, default=None, help='name for network in generated reports') parser.add_argument('--batch_size', type=int, default=25, help='batch size for forward pass') parser.add_argument('--mixed_units', action='store_true', default=False, help='true to keep alpha for non-zeroed units') parser.add_argument('--size', type=int, default=200, help='number of images to test') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA usage') parser.add_argument('--quiet', action='store_true', default=False, help='silences console output') if len(sys.argv) == 1: parser.print_usage(sys.stderr) sys.exit(1) args = parser.parse_args() # Set up console output verbose_progress(not args.quiet) # Speed up pytorch torch.backends.cudnn.benchmark = True # Set up CUDA args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: torch.backends.cudnn.benchmark = True # Take defaults for model constructor etc from dissect.json settings. with open(os.path.join(args.outdir, 'dissect.json')) as f: dissection = EasyDict(json.load(f)) if args.model is None: args.model = dissection.settings.model if args.pthfile is None: args.pthfile = dissection.settings.pthfile if args.segmenter is None: args.segmenter = dissection.settings.segmenter if args.layer is None: args.layer = dissection.settings.layers[0] args.layers = [args.layer] # Also load specific analysis layername = args.layer[1] if args.metric == 'iou': summary = dissection else: with open(os.path.join(args.outdir, layername, args.metric, args.classname, 'summary.json')) as f: summary = EasyDict(json.load(f)) # Instantiate generator model = create_instrumented_model(args, gen=True, edit=True) if model is None: print('No model specified') sys.exit(1) # Instantiate model device = next(model.parameters()).device input_shape = model.input_shape # 4d input if convolutional, 2d input if first layer is linear. raw_sample = standard_z_sample(args.size, input_shape[1], seed=3).view( (args.size,) + input_shape[1:]) dataset = TensorDataset(raw_sample) # Create the segmenter segmenter = autoimport_eval(args.segmenter) # Now do the actual work. labelnames, catnames = ( segmenter.get_label_and_category_names(dataset)) label_category = [catnames.index(c) if c in catnames else 0 for l, c in labelnames] labelnum_from_name = {n[0]: i for i, n in enumerate(labelnames)} segloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, num_workers=10, pin_memory=(device.type == 'cuda')) # Index the dissection layers by layer name. # First, collect a baseline for l in model.ablation: model.ablation[l] = None # For each sort-order, do an ablation progress = default_progress() classname = args.classname classnum = labelnum_from_name[classname] # Get iou ranking from dissect.json iou_rankname = '%s-%s' % (classname, 'iou') dissect_layer = {lrec.layer: lrec for lrec in dissection.layers} iou_ranking = next(r for r in dissect_layer[layername].rankings if r.name == iou_rankname) # Get trained ranking from summary.json rankname = '%s-%s' % (classname, args.metric) summary_layer = {lrec.layer: lrec for lrec in summary.layers} ranking = next(r for r in summary_layer[layername].rankings if r.name == rankname) # Get ordering, first by ranking, then break ties by iou. ordering = [t[2] for t in sorted([(s1, s2, i) for i, (s1, s2) in enumerate(zip(ranking.score, iou_ranking.score))])] values = (-numpy.array(ranking.score))[ordering] if not args.mixed_units: values[...] = 1 ablationdir = os.path.join(args.outdir, layername, 'fullablation') measurements = measure_full_ablation(segmenter, segloader, model, classnum, layername, ordering[:args.unitcount], values[:args.unitcount]) measurements = measurements.cpu().numpy().tolist() os.makedirs(ablationdir, exist_ok=True) with open(os.path.join(ablationdir, '%s.json'%rankname), 'w') as f: json.dump(dict( classname=classname, classnum=classnum, baseline=measurements[0], layer=layername, metric=args.metric, ablation_units=ordering, ablation_values=values.tolist(), ablation_effects=measurements[1:]), f)