def create_instrumented_model(args, **kwargs): ''' Creates an instrumented model out of a namespace of arguments that correspond to ArgumentParser command-line args: model: a string to evaluate as a constructor for the model. pthfile: (optional) filename of .pth file for the model. layers: a list of layers to instrument, defaulted if not provided. edit: True to instrument the layers for editing. gen: True for a generator model. One-pixel input assumed. imgsize: For non-generator models, (y, x) dimensions for RGB input. cuda: True to use CUDA. The constructed model will be decorated with the following attributes: input_shape: (usually 4d) tensor shape for single-image input. output_shape: 4d tensor shape for output. feature_shape: map of layer names to 4d tensor shape for featuremaps. retained: map of layernames to tensors, filled after every evaluation. ablation: if editing, map of layernames to [0..1] alpha values to fill. replacement: if editing, map of layernames to values to fill. When editing, the feature value x will be replaced by: `x = (replacement * ablation) + (x * (1 - ablation))` ''' args = EasyDict(vars(args), **kwargs) # Construct the network if args.model is None: print_progress('No model specified') return None if isinstance(args.model, torch.nn.Module): model = args.model else: model = autoimport_eval(args.model) # Unwrap any DataParallel-wrapped model if isinstance(model, torch.nn.DataParallel): model = next(model.children()) # Load its state dict meta = {} if getattr(args, 'pthfile', None) 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) model.meta = meta # Decide which layers to instrument. if getattr(args, 'layer', None) is not None: args.layers = [args.layer] if getattr(args, 'layers', None) is None: # Skip wrappers with only one named model container = model prefix = '' while len(list(container.named_children())) == 1: name, container = next(container.named_children()) prefix += name + '.' # Default to all nontrivial top-level layers except last. args.layers = [ prefix + name for name, module in container.named_children() if type(module).__module__ not in [ # Skip ReLU and other activations. 'torch.nn.modules.activation', # Skip pooling layers. 'torch.nn.modules.pooling' ] ][:-1] print_progress('Defaulting to layers: %s' % ' '.join(args.layers)) # Instrument the layers. retain_layers(model, args.layers) if getattr(args, 'edit', False): edit_layers(model, args.layers) model.eval() if args.cuda: model.cuda() # Annotate input, output, and feature shapes annotate_model_shapes(model, gen=getattr(args, 'gen', False), imgsize=getattr(args, 'imgsize', None)) return model
def main(): # Training settings def strpair(arg): p = tuple(arg.split(':')) if len(p) == 1: p = p + p return p def intpair(arg): p = arg.split(',') if len(p) == 1: p = p + p return tuple(int(v) for v in p) parser = argparse.ArgumentParser( description='Net dissect utility', prog='python -m netdissect', 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', help='directory for dissection output') parser.add_argument('--layers', type=strpair, nargs='+', help='space-separated list of layer names to dissect' + ', in the form layername[:reportedname]') parser.add_argument('--segments', type=str, default='dataset/broden', help='directory containing segmentation dataset') parser.add_argument('--download', action='store_true', default=False, help='downloads Broden dataset if needed') parser.add_argument('--imgsize', type=intpair, default=(227, 227), help='input image size to use') parser.add_argument('--netname', type=str, default=None, help='name for network in generated reports') parser.add_argument('--meta', type=str, nargs='+', help='json files of metadata to add to report') parser.add_argument('--examples', type=int, default=20, help='number of image examples per unit') parser.add_argument('--size', type=int, default=10000, help='dataset subset size to use') parser.add_argument('--broden_version', type=int, default=1, help='broden version number') parser.add_argument('--batch_size', type=int, default=100, help='batch size for forward pass') parser.add_argument('--num_workers', type=int, default=24, help='number of DataLoader workers') parser.add_argument('--quantile_threshold', type=float, default=None, help='quantile to use for masks') parser.add_argument('--no-labels', action='store_true', default=False, help='disables labeling of units') parser.add_argument('--ablation', action='store_true', default=False, help='enables single unit ablation of units') parser.add_argument('--iqr', action='store_true', default=False, help='enables iqr calculation') parser.add_argument('--maxiou', action='store_true', default=False, help='enables maxiou calculation') parser.add_argument('--covariance', action='store_true', default=False, help='enables covariance calculation') parser.add_argument('--no-images', action='store_true', default=False, help='disables generation of unit images') parser.add_argument('--single-images', action='store_true', default=False, help='generates single images also') parser.add_argument('--no-report', action='store_true', default=False, help='disables generation report summary') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA usage') parser.add_argument('--gan', type=str, default=None, help='netdissect.GanImageSegmenter() to probe a GAN') parser.add_argument('--perturbation', default=None, help='filename of perturbation attack to apply') parser.add_argument('--add_scale_offset', action='store_true', default=None, help='offsets masks according to stride and padding') 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() args.images = not args.no_images args.report = not args.no_report args.labels = not args.no_labels # Set up console output verbose_progress(not args.quiet) # Speed up pytorch torch.backends.cudnn.benchmark = True # Special case: download flag without model to test. if args.model is None and args.download: from netdissect.broden import ensure_broden_downloaded for resolution in [224, 227, 384]: ensure_broden_downloaded(args.segments, resolution, args.broden_version) sys.exit(0) # Help if broden is not present if not os.path.isdir(args.segments): print_progress('Segmentation dataset not found at %s.' % args.segments) print_progress('Specify dataset directory using --segments [DIR]') print_progress('To download Broden, run: netdissect --download') sys.exit(1) # Default threshold if args.quantile_threshold is None: if args.gan: args.quantile_threshold = 0.01 else: args.quantile_threshold = 0.005 # Construct the network if args.model is None: print_progress('No model specified') sys.exit(1) # Set up CUDA args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: torch.backends.cudnn.benchmark = True model = autoimport_eval(args.model) # Unwrap any DataParallel-wrapped model if isinstance(model, torch.nn.DataParallel): model = next(model.children()) # Default add_scale_offset only for AlexNet-looking models. if args.add_scale_offset is None and not args.gan: args.add_scale_offset = ('Alex' in model.__class__.__name__) # Load its state dict meta = {} if args.pthfile is None: print_progress('Dissecting model without pth file.') else: 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) # Update any metadata from files, if any if args.meta: for mfilename in args.meta: with open(mfilename) as f: meta.update(json.load(f)) # Instrument it and prepare it for eval if not args.layers: # Skip wrappers with only one named modele container = model prefix = '' while len(list(container.named_children())) == 1: name, container = next(container.named_children()) prefix += name + '.' # Default to all nontrivial top-level layers except last. args.layers = [ prefix + name for name, module in container.named_children() if type(module).__module__ not in [ # Skip ReLU and other activations. 'torch.nn.modules.activation', # Skip pooling layers. 'torch.nn.modules.pooling' ] ][:-1] print_progress('Defaulting to layers: %s' % ' '.join(args.layers)) retain_layers(model, args.layers, args.add_scale_offset) if args.gan: ablate_layers(model, args.layers) model.eval() if args.cuda: model.cuda() # Set up the output directory, verify write access if args.outdir is None: args.outdir = os.path.join('dissect', type(model).__name__) print_progress('Writing output into %s.' % args.outdir) os.makedirs(args.outdir, exist_ok=True) train_dataset = None if not args.gan: # Load dataset for ordinary case. # Load perturbation perturbation = numpy.load( args.perturbation) if args.perturbation else None # Load broden dataset dataset = try_to_load_broden(args.segments, args.imgsize, args.broden_version, perturbation, args.download, args.size) if dataset is None: ds = try_to_load_multiseg(args.segments, args.imgsize, perturbation, args.size) if dataset is None: print_progress('No segmentation dataset found in %s' % args.segements) print_progress('use --download to download Broden.') sys.exit(1) recovery = ReverseNormalize(IMAGE_MEAN, IMAGE_STDEV) else: # 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)): sample = standard_z_sample(args.size, first_layer.in_channels)[:, :, None, None] train_sample = standard_z_sample(args.size, first_layer.in_channels, seed=2)[:, :, None, None] else: sample = standard_z_sample(args.size, first_layer.in_features) train_sample = standard_z_sample(args.size, first_layer.in_features, seed=2) dataset = TensorDataset(sample) train_dataset = TensorDataset(train_sample) recovery = autoimport_eval(args.gan) # Run dissect dissect(args.outdir, model, dataset, train_dataset=train_dataset, recover_image=recovery, examples_per_unit=args.examples, netname=args.netname, quantile_threshold=args.quantile_threshold, meta=meta, make_images=args.images or args.single_images, make_labels=args.labels, make_ablation=args.ablation, make_iqr=args.iqr, make_maxiou=args.maxiou, make_covariance=args.covariance, make_report=args.report, make_row_images=args.images, make_single_images=args.single_images, batch_size=args.batch_size, num_workers=args.num_workers, settings=vars(args))
def eval_loss_and_reg(): discrete_experiments = dict( # dpixel=dict(discrete_pixels=True), # dunits20=dict(discrete_units=20), # dumix20=dict(discrete_units=20, mixed_units=True), # dunits10=dict(discrete_units=10), # abonly=dict(ablation_only=True), # fimabl=dict(ablation_only=True, # fullimage_ablation=True, # fullimage_measurement=True), dboth20=dict(discrete_units=20, discrete_pixels=True), # dbothm20=dict(discrete_units=20, mixed_units=True, # discrete_pixels=True), # abdisc20=dict(discrete_units=20, discrete_pixels=True, # ablation_only=True), # abdiscm20=dict(discrete_units=20, mixed_units=True, # discrete_pixels=True, # ablation_only=True), # fimadp=dict(discrete_pixels=True, # ablation_only=True, # fullimage_ablation=True, # fullimage_measurement=True), # fimadu10=dict(discrete_units=10, # ablation_only=True, # fullimage_ablation=True, # fullimage_measurement=True), # fimadb10=dict(discrete_units=10, discrete_pixels=True, # ablation_only=True, # fullimage_ablation=True, # fullimage_measurement=True), fimadbm10=dict(discrete_units=10, mixed_units=True, discrete_pixels=True, ablation_only=True, fullimage_ablation=True, fullimage_measurement=True), # fimadu20=dict(discrete_units=20, # ablation_only=True, # fullimage_ablation=True, # fullimage_measurement=True), # fimadb20=dict(discrete_units=20, discrete_pixels=True, # ablation_only=True, # fullimage_ablation=True, # fullimage_measurement=True), fimadbm20=dict(discrete_units=20, mixed_units=True, discrete_pixels=True, ablation_only=True, fullimage_ablation=True, fullimage_measurement=True) ) with torch.no_grad(): total_loss = 0 discrete_losses = {k: 0 for k in discrete_experiments} for [pbatch, ploc, cbatch, cloc] in progress( torch.utils.data.DataLoader(TensorDataset( corpus.eval_present_sample, corpus.eval_present_location, corpus.eval_candidate_sample, corpus.eval_candidate_location), batch_size=args.inference_batch_size, num_workers=10, shuffle=False, pin_memory=True), desc="Eval"): # First, put in zeros for the selected units. # Loss is amount of remaining object. total_loss = total_loss + ace_loss(segmenter, classnum, model, args.layer, high_replacement, ablation, pbatch, ploc, cbatch, cloc, run_backward=False, ablation_only=ablation_only, fullimage_measurement=fullimage_measurement) for k, config in discrete_experiments.items(): discrete_losses[k] = discrete_losses[k] + ace_loss( segmenter, classnum, model, args.layer, high_replacement, ablation, pbatch, ploc, cbatch, cloc, run_backward=False, **config) avg_loss = (total_loss / args.eval_size).item() avg_d_losses = {k: (d / args.eval_size).item() for k, d in discrete_losses.items()} regularizer = (args.l2_lambda * ablation.pow(2).sum()) print_progress('Epoch %d Loss %g Regularizer %g' % (epoch, avg_loss, regularizer)) print_progress(' '.join('%s: %g' % (k, d) for k, d in avg_d_losses.items())) print_progress(scale_summary(ablation.view(-1), 10, 3)) return avg_loss, regularizer, avg_d_losses
def main(): # Training settings def strpair(arg): p = tuple(arg.split(':')) if len(p) == 1: p = p + p return p def intpair(arg): p = arg.split(',') if len(p) == 1: p = p + p return tuple(int(v) for v in p) parser = argparse.ArgumentParser( description='Net dissect utility', prog='python -m netdissect', 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('--unstrict', action='store_true', default=False, help='ignore unexpected pth parameters') parser.add_argument('--submodule', type=str, default=None, help='submodule to load from pthfile') parser.add_argument('--outdir', type=str, default='dissect', help='directory for dissection output') parser.add_argument('--layers', type=strpair, nargs='+', help='space-separated list of layer names to dissect' + ', in the form layername[:reportedname]') parser.add_argument('--segments', type=str, default='dataset/broden', help='directory containing segmentation dataset') parser.add_argument('--segmenter', type=str, default=None, help='constructor for asegmenter class') parser.add_argument('--download', action='store_true', default=False, help='downloads Broden dataset if needed') parser.add_argument('--imagedir', type=str, default=None, help='directory containing image-only dataset') parser.add_argument('--imgsize', type=intpair, default=(227, 227), help='input image size to use') parser.add_argument('--netname', type=str, default=None, help='name for network in generated reports') parser.add_argument('--meta', type=str, nargs='+', help='json files of metadata to add to report') parser.add_argument('--merge', type=str, help='json file of unit data to merge in report') parser.add_argument('--examples', type=int, default=20, help='number of image examples per unit') parser.add_argument('--size', type=int, default=10000, help='dataset subset size to use') parser.add_argument('--batch_size', type=int, default=100, help='batch size for forward pass') parser.add_argument('--num_workers', type=int, default=24, help='number of DataLoader workers') parser.add_argument('--quantile_threshold', type=strfloat, default=None, choices=[FloatRange(0.0, 1.0), 'iqr'], help='quantile to use for masks') parser.add_argument('--no-labels', action='store_true', default=False, help='disables labeling of units') parser.add_argument('--maxiou', action='store_true', default=False, help='enables maxiou calculation') parser.add_argument('--covariance', action='store_true', default=False, help='enables covariance calculation') parser.add_argument('--rank_all_labels', action='store_true', default=False, help='include low-information labels in rankings') parser.add_argument('--no-images', action='store_true', default=False, help='disables generation of unit images') parser.add_argument('--no-report', action='store_true', default=False, help='disables generation report summary') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA usage') parser.add_argument('--gen', action='store_true', default=False, help='test a generator model (e.g., a GAN)') parser.add_argument('--gan', action='store_true', default=False, help='synonym for --gen') parser.add_argument('--perturbation', default=None, help='filename of perturbation attack to apply') parser.add_argument('--add_scale_offset', action='store_true', default=None, help='offsets masks according to stride and padding') 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() args.images = not args.no_images args.report = not args.no_report args.labels = not args.no_labels if args.gan: args.gen = args.gan # Set up console output verbose_progress(not args.quiet) # Exit right away if job is already done or being done. if args.outdir is not None: exit_if_job_done(args.outdir) # Speed up pytorch torch.backends.cudnn.benchmark = True # Special case: download flag without model to test. if args.model is None and args.download: from netdissect.broden import ensure_broden_downloaded for resolution in [224, 227, 384]: ensure_broden_downloaded(args.segments, resolution, 1) from netdissect.segmenter import ensure_upp_segmenter_downloaded ensure_upp_segmenter_downloaded('dataset/segmodel') sys.exit(0) # Help if broden is not present if not args.gen and not args.imagedir and not os.path.isdir(args.segments): print_progress('Segmentation dataset not found at %s.' % args.segments) print_progress('Specify dataset directory using --segments [DIR]') print_progress('To download Broden, run: netdissect --download') sys.exit(1) # Default segmenter class if args.gen and args.segmenter is None: args.segmenter = ("netdissect.segmenter.UnifiedParsingSegmenter(" + "segsizes=[256], segdiv=None)") # Default threshold if args.quantile_threshold is None: if args.gen: args.quantile_threshold = 'iqr' else: args.quantile_threshold = 0.005 # Set up CUDA args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: torch.backends.cudnn.benchmark = True # Construct the network with specified layers instrumented if args.model is None: print_progress('No model specified') sys.exit(1) model = create_instrumented_model(args) # Update any metadata from files, if any meta = getattr(model, 'meta', {}) if args.meta: for mfilename in args.meta: with open(mfilename) as f: meta.update(json.load(f)) # Load any merge data from files mergedata = None if args.merge: with open(args.merge) as f: mergedata = json.load(f) # Set up the output directory, verify write access if args.outdir is None: args.outdir = os.path.join('dissect', type(model).__name__) exit_if_job_done(args.outdir) print_progress('Writing output into %s.' % args.outdir) os.makedirs(args.outdir, exist_ok=True) train_dataset = None if not args.gen: # Load dataset for classifier case. # Load perturbation perturbation = numpy.load( args.perturbation) if args.perturbation else None segrunner = None # Load broden dataset if args.imagedir is not None: dataset = try_to_load_images(args.imagedir, args.imgsize, perturbation, args.size) segrunner = ImageOnlySegRunner(dataset) else: dataset = try_to_load_broden(args.segments, args.imgsize, 1, perturbation, args.download, args.size) if dataset is None: dataset = try_to_load_multiseg(args.segments, args.imgsize, perturbation, args.size) if dataset is None: print_progress('No segmentation dataset found in %s', args.segments) print_progress('use --download to download Broden.') sys.exit(1) else: # For segmenter case the dataset is just a random z #dataset = z_dataset_for_model(model, args.size) #train_dataset = z_dataset_for_model(model, args.size, seed=2) dataset = datasets.ImageFolder('dataset/Adissect', transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize( (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) train_dataset = dataset segrunner = GeneratorSegRunner(autoimport_eval(args.segmenter)) torch.no_grad() # Run dissect dissect(args.outdir, model, dataset, train_dataset=train_dataset, segrunner=segrunner, examples_per_unit=args.examples, netname=args.netname, quantile_threshold=args.quantile_threshold, meta=meta, merge=mergedata, make_images=args.images, make_labels=args.labels, make_maxiou=args.maxiou, make_covariance=args.covariance, make_report=args.report, make_row_images=args.images, make_single_images=True, rank_all_labels=args.rank_all_labels, batch_size=args.batch_size, num_workers=args.num_workers, settings=vars(args)) # Mark the directory so that it's not done again. mark_job_done(args.outdir)
def main(): # Training settings def strpair(arg): p = tuple(arg.split(':')) if len(p) == 1: p = p + p return p def intpair(arg): p = arg.split(',') if len(p) == 1: p = p + p return tuple(int(v) for v in p) parser = argparse.ArgumentParser( description='Net dissect utility', prog='python -m netdissect', 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', help='directory for dissection output') parser.add_argument('--layers', type=strpair, nargs='+', help='space-separated list of layer names to dissect' + ', in the form layername[:reportedname]') parser.add_argument('--segments', type=str, default='dataset/broden', help='directory containing segmentation dataset') parser.add_argument('--download', action='store_true', default=False, help='downloads Broden dataset if needed') parser.add_argument('--imgsize', type=intpair, default=(227, 227), help='input image size to use') parser.add_argument('--netname', type=str, default=None, help='name for network in generated reports') parser.add_argument('--meta', type=str, nargs='+', help='json files of metadata to add to report') parser.add_argument('--examples', type=int, default=20, help='number of image examples per unit') parser.add_argument('--size', type=int, default=10000, help='dataset subset size to use') parser.add_argument('--broden_version', type=int, default=1, help='broden version number') parser.add_argument('--batch_size', type=int, default=100, help='batch size for forward pass') parser.add_argument('--num_workers', type=int, default=24, help='number of DataLoader workers') parser.add_argument('--no-labels', action='store_true', default=False, help='disables labeling of units') parser.add_argument('--no-images', action='store_true', default=False, help='disables generation of unit images') parser.add_argument('--single-images', action='store_true', default=False, help='generates single images also') parser.add_argument('--no-report', action='store_true', default=False, help='disables generation report summary') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA usage') parser.add_argument('--perturbation', default=None, help='filename of perturbation attack to apply') parser.add_argument('--add_scale_offset', action='store_true', default=None, help='offsets masks according to stride and padding') 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() args.images = not args.no_images args.report = not args.no_report args.labels = not args.no_labels # Set up console output verbose_progress(not args.quiet) # Speed up pytorch torch.backends.cudnn.benchmark = True # Special case: download flag without model to test. if args.model is None and args.download: from netdissect.broden import ensure_broden_downloaded for resolution in [224, 227, 384]: ensure_broden_downloaded(args.segments, resolution, args.broden_version) sys.exit(0) # Help if broden is not present if not os.path.isdir(args.segments): print_progress('Segmentation dataset not found at %s.' % args.segments) print_progress('Specify dataset directory using --segments [DIR]') print_progress('To download Broden, run: netdissect --download') sys.exit(1) # Construct the network if args.model is None: print_progress('No model specified') sys.exit(1) # Set up CUDA args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: torch.backends.cudnn.benchmark = True model = eval_constructor(args.model) # Default add_scale_offset only for AlexNet-looking models. if args.add_scale_offset is None: args.add_scale_offset = ('Alex' in model.__class__.__name__) # Load its state dict meta = {} if args.pthfile is None: print_progress('Dissecting model without pth file.') else: 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) # Update any metadata from files, if any if args.meta: for mfilename in args.meta: with open(mfilename) as f: meta.update(json.load(f)) # Instrument it and prepare it for eval if not args.layers: print_progress('No layers specified') sys.exit(1) retain_layers(model, args.layers, args.add_scale_offset) model.eval() if args.cuda: model.cuda() # Set up the output directory, verify write access if args.outdir is None: args.outdir = os.path.join('dissect', type(model).__name__) print_progress('Writing output into %s.' % args.outdir) os.makedirs(args.outdir, exist_ok=True) # Load perturbation perturbation = numpy.load(args.perturbation) if args.perturbation else None # Load broden dataset ds = try_to_load_broden(args.segments, args.imgsize, args.broden_version, perturbation, args.download, args.size) if ds is None: ds = try_to_load_multiseg(args.segments, args.imgsize, perturbation, args.size) if ds is None: print_progress('No segmentation dataset found in %s' % args.segements) print_progress('use --download to download Broden.') sys.exit(1) # Run dissect dissect(args.outdir, model, ds, recover_image=ReverseNormalize(IMAGE_MEAN, IMAGE_STDEV), examples_per_unit=args.examples, netname=args.netname, meta=meta, make_images=args.images, make_labels=args.labels, make_report=args.report, make_single_images=args.single_images, batch_size=args.batch_size, num_workers=args.num_workers, settings=vars(args))
def create_instrumented_model(args, **kwargs): ''' Creates an instrumented model out of a namespace of arguments that correspond to ArgumentParser command-line args: model: a string to evaluate as a constructor for the model. pthfile: (optional) filename of .pth file for the model. layers: a list of layers to instrument, defaulted if not provided. edit: True to instrument the layers for editing. gen: True for a generator model. One-pixel input assumed. imgsize: For non-generator models, (y, x) dimensions for RGB input. cuda: True to use CUDA. The constructed model will be decorated with the following attributes: input_shape: (usually 4d) tensor shape for single-image input. output_shape: 4d tensor shape for output. feature_shape: map of layer names to 4d tensor shape for featuremaps. retained: map of layernames to tensors, filled after every evaluation. ablation: if editing, map of layernames to [0..1] alpha values to fill. replacement: if editing, map of layernames to values to fill. When editing, the feature value x will be replaced by: `x = (replacement * ablation) + (x * (1 - ablation))` ''' args = EasyDict(vars(args), **kwargs) # Construct the network if args.model is None: print_progress('No model specified') return None #if isinstance(args.model, torch.nn.Module): # model = args.model #else: # model = autoimport_eval(args.model) model = UnetNormalized() checkpoint = torch.load('p2p_churches.pth') model.load_state_dict(checkpoint) model.cuda() #model.eval() torch.no_grad() # Unwrap any DataParallel-wrapped model if isinstance(model, torch.nn.DataParallel): model = next(model.children()) # Load its state dict #meta = {} #if getattr(args, 'pthfile', None) 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'] # submodule = getattr(args, 'submodule', None) # if submodule is not None and len(submodule): # remove_prefix = submodule + '.' # data = { k[len(remove_prefix):]: v for k, v in data.items() # if k.startswith(remove_prefix)} # if not len(data): # print_progress('No submodule %s found in %s' % # (submodule, args.pthfile)) # return None # model.load_state_dict(data, strict=not getattr(args, 'unstrict', False)) # Decide which layers to instrument. #if getattr(args, 'layer', None) is not None: # args.layers = [args.layer] args.layers = [layer5, layer9, layer12] #if getattr(args, 'layers', None) is None: # # Skip wrappers with only one named model # container = model # prefix = '' # while len(list(container.named_children())) == 1: # name, container = next(container.named_children()) # prefix += name + '.' # # Default to all nontrivial top-level layers except last. # args.layers = [prefix + name # for name, module in container.named_children() # if type(module).__module__ not in [ # # Skip ReLU and other activations. # 'torch.nn.modules.activation', # # Skip pooling layers. # 'torch.nn.modules.pooling'] # ][:-1] # print_progress('Defaulting to layers: %s' % ' '.join(args.layers)) # Now wrap the model for instrumentation. model = InstrumentedModel(model) # model.meta = meta # Instrument the layers. model.retain_layers(args.layers) #model.eval() #if args.cuda: # model.cuda() # Annotate input, output, and feature shapes annotate_model_shapes(model, gen=getattr(args, 'gen', False), imgsize=getattr(args, 'imgsize', None)) return model
def train_ablation(args, corpus, cachefile, model, segmenter, classnum, initial_ablation=None): progress = default_progress() cachedir = os.path.dirname(cachefile) snapdir = os.path.join(cachedir, 'snapshots') os.makedirs(snapdir, exist_ok=True) # high_replacement = corpus.feature_99[None,:,None,None].cuda() high_replacement = (corpus.weighted_mean_present_feature[None, :, None, None].cuda()) high_replacement.requires_grad = False for p in model.parameters(): p.requires_grad = False ablation = torch.zeros(high_replacement.shape).cuda() if initial_ablation is not None: ablation.view(-1)[...] = initial_ablation ablation.requires_grad = True optimizer = torch.optim.Adam([ablation], lr=0.01) start_epoch = 0 if not args.no_cache: for start_epoch in reversed(range(args.train_epochs)): snapfile = os.path.join(snapdir, 'epoch-%d.pth' % start_epoch) if os.path.exists(snapfile): data = torch.load(snapfile) with torch.no_grad(): ablation[...] = data['ablation'].to(ablation.device) optimizer.load_state_dict(data['optimizer']) start_epoch += 1 break update_size = args.train_update_freq * args.train_batch_size for epoch in range(start_epoch, args.train_epochs): candidate_shuffle = torch.randperm(len(corpus.candidate_sample)) train_loss = 0 for batch_num, [pbatch, ploc, cbatch, cloc] in enumerate( progress(torch.utils.data.DataLoader( TensorDataset( corpus.object_present_sample, corpus.object_present_location, corpus.candidate_sample[candidate_shuffle], corpus.candidate_location[candidate_shuffle]), batch_size=args.train_batch_size, num_workers=10, shuffle=True, pin_memory=True), desc="ACE opt epoch %d" % epoch)): if batch_num % args.train_update_freq == 0: optimizer.zero_grad() # First, put in zeros for the selected units. Loss is amount # of remaining object. loss = ace_loss(segmenter, classnum, model, args.layer, high_replacement, ablation, pbatch, ploc, cbatch, cloc, run_backward=True) with torch.no_grad(): train_loss = train_loss + loss if (batch_num + 1) % args.train_update_freq == 0: # Third, add some L2 loss to encourage sparsity. regularizer = (args.l2_lambda * update_size * ablation.pow(2).sum()) regularizer.backward() optimizer.step() with torch.no_grad(): ablation.clamp_(0, 1) post_progress(l=(train_loss / update_size).item(), r=(regularizer / update_size).item()) train_loss = 0 with torch.no_grad(): total_loss = 0 for [pbatch, ploc, cbatch, cloc] in progress(torch.utils.data.DataLoader( TensorDataset(corpus.eval_present_sample, corpus.eval_present_location, corpus.eval_candidate_sample, corpus.eval_candidate_location), batch_size=args.inference_batch_size, num_workers=10, shuffle=False, pin_memory=True), desc="Eval"): # First, put in zeros for the selected units. Loss is amount # of remaining object. total_loss = total_loss + ace_loss(segmenter, classnum, model, args.layer, high_replacement, ablation, pbatch, ploc, cbatch, cloc, run_backward=False) avg_loss = (total_loss / args.eval_size).item() regularizer = (args.l2_lambda * ablation.pow(2).sum()) print_progress('Epoch %d Loss %g, Regularizer %g' % (epoch, avg_loss, regularizer)) print_progress(scale_summary(ablation.view(-1), 10, 3)) torch.save( dict(ablation=ablation, optimizer=optimizer.state_dict(), avg_loss=avg_loss), os.path.join(snapdir, 'epoch-%d.pth' % epoch)) numpy.save(os.path.join(snapdir, 'epoch-%d.npy' % epoch), ablation.detach().cpu().numpy()) # The output of this phase is this set of scores. return ablation.view(-1).detach().cpu().numpy()
def main(): # Training settings def strpair(arg): p = tuple(arg.split(':')) if len(p) == 1: p = p + p return p parser = argparse.ArgumentParser( description='Net dissect utility', 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', 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('--startcount', type=int, default=1, help='number of units to ablate') parser.add_argument('--unitcount', type=int, default=30, help='number of units to ablate') parser.add_argument('--segmenter', type=str, default='dataset/broden', 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=1000, 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 # Construct the network if args.model is None: print_progress('No model specified') sys.exit(1) # Set up CUDA args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: torch.backends.cudnn.benchmark = True model = autoimport_eval(args.model) # Unwrap any DataParallel-wrapped model if isinstance(model, torch.nn.DataParallel): model = next(model.children()) # Load its state dict meta = {} if args.pthfile is None: print_progress('Dissecting model without pth file.') else: 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) # Instrument it and prepare it for eval if not args.layers: # Skip wrappers with only one named modele container = model prefix = '' while len(list(container.named_children())) == 1: name, container = next(container.named_children()) prefix += name + '.' # Default to all nontrivial top-level layers except last. args.layers = [ prefix + name for name, module in container.named_children() if type(module).__module__ not in [ # Skip ReLU and other activations. 'torch.nn.modules.activation', # Skip pooling layers. 'torch.nn.modules.pooling' ] ][:-1] print_progress('Defaulting to layers: %s' % ' '.join(args.layers)) edit_layers(model, args.layers) model.eval() if args.cuda: model.cuda() # Set up the output directory, verify write access if args.outdir is None: args.outdir = os.path.join('dissect', type(model).__name__) print_progress('Writing output into %s.' % args.outdir) os.makedirs(args.outdir, exist_ok=True) train_dataset = None # 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)): sample = standard_z_sample(args.size, first_layer.in_channels)[:, :, None, None] train_sample = standard_z_sample(args.size, first_layer.in_channels, seed=2)[:, :, None, None] else: sample = standard_z_sample(args.size, first_layer.in_features) train_sample = standard_z_sample(args.size, first_layer.in_features, seed=2) dataset = TensorDataset(sample) train_dataset = TensorDataset(train_sample) recovery = autoimport_eval(args.segmenter) # Now do the actual work. device = next(model.parameters()).device labelnames, catnames = (recovery.get_label_and_category_names(dataset)) label_category = [catnames.index(c) 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')) with open(os.path.join(args.outdir, 'dissect.json'), 'r') as f: dissect = EasyDict(json.load(f)) # Index the dissection layers by layer name. dissect_layer = {lrec.layer: lrec for lrec in dissect.layers} # First, collect a baseline for l in model.ablation: model.ablation[l] = None baseline = count_segments(recovery, segloader, model) # For each sort-order, do an ablation progress = default_progress() for classname in progress(args.classes): post_progress(c=classname) for layername in progress(model.ablation): post_progress(l=layername) rankname = '%s-%s' % (classname, args.metric) measurements = {} 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, 'ablation') 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)): import pdb pdb.set_trace() continue if len(data.ablation_effects) >= args.unitcount: continue # file already done. measurements = data.ablation_effects for count in progress(range(args.startcount, min(args.unitcount, len(ordering)) + 1), desc='units'): if str(count) in measurements: continue ablation = numpy.zeros(len(ranking.score), dtype='float32') ablation[ordering[:count]] = 1 for l in model.ablation: model.ablation[l] = ablation if layername == l else None m = count_segments(recovery, segloader, model)[classnum].item() print_progress( '%s %s %d units (#%d), %g -> %g' % (layername, rankname, count, ordering[count - 1].item(), baseline[classnum].item(), m)) measurements[str(count)] = m 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=baseline[classnum].item(), layer=layername, metric=args.metric, ablation_units=ordering.tolist(), ablation_effects=measurements), f)