Spaces:
Build error
Build error
| import torch, sys, os, argparse, textwrap, numbers, numpy, json, PIL | |
| from torchvision import transforms | |
| from torch.utils.data import TensorDataset | |
| from netdissect.progress import default_progress, post_progress, desc_progress | |
| from netdissect.progress import verbose_progress, print_progress | |
| from netdissect.nethook import edit_layers | |
| from netdissect.zdataset import standard_z_sample | |
| from netdissect.autoeval import autoimport_eval | |
| from netdissect.easydict import EasyDict | |
| from netdissect.modelconfig import create_instrumented_model | |
| help_epilog = '''\ | |
| Example: | |
| python -m netdissect.evalablate \ | |
| --segmenter "netdissect.segmenter.UnifiedParsingSegmenter(segsizes=[256], segdiv='quad')" \ | |
| --model "proggan.from_pth_file('models/lsun_models/${SCENE}_lsun.pth')" \ | |
| --outdir dissect/dissectdir \ | |
| --classes mirror coffeetable tree \ | |
| --layers layer4 \ | |
| --size 1000 | |
| Output layout: | |
| dissectdir/layer5/ablation/mirror-iqr.json | |
| { class: "mirror", | |
| classnum: 43, | |
| pixel_total: 41342300, | |
| class_pixels: 1234531, | |
| layer: "layer5", | |
| ranking: "mirror-iqr", | |
| ablation_units: [341, 23, 12, 142, 83, ...] | |
| ablation_pixels: [143242, 132344, 429931, ...] | |
| } | |
| ''' | |
| 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 | |
| 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 | |
| # 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 | |
| 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) | |
| 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 measure_ablation(segmenter, loader, model, classnum, layer, ordering): | |
| total_bincount = 0 | |
| data_size = 0 | |
| device = next(model.parameters()).device | |
| progress = default_progress() | |
| for l in model.ablation: | |
| model.ablation[l] = None | |
| feature_units = model.feature_shape[layer][1] | |
| feature_shape = model.feature_shape[layer][2:] | |
| repeats = len(ordering) | |
| total_scores = torch.zeros(repeats + 1) | |
| for i, batch in enumerate(progress(loader)): | |
| z_batch = batch[0] | |
| model.ablation[layer] = None | |
| tensor_images = model(z_batch.to(device)) | |
| seg = segmenter.segment_batch(tensor_images, downsample=2) | |
| mask = (seg == classnum).max(1)[0] | |
| downsampled_seg = torch.nn.functional.adaptive_avg_pool2d( | |
| mask.float()[:,None,:,:], feature_shape)[:,0,:,:] | |
| total_scores[0] += downsampled_seg.sum().cpu() | |
| # Now we need to do an intervention for every location | |
| # that had a nonzero downsampled_seg, if any. | |
| interventions_needed = downsampled_seg.nonzero() | |
| location_count = len(interventions_needed) | |
| if location_count == 0: | |
| continue | |
| interventions_needed = interventions_needed.repeat(repeats, 1) | |
| inter_z = batch[0][interventions_needed[:,0]].to(device) | |
| inter_chan = torch.zeros(repeats, location_count, feature_units, | |
| device=device) | |
| for j, u in enumerate(ordering): | |
| inter_chan[j:, :, u] = 1 | |
| inter_chan = inter_chan.view(len(inter_z), feature_units) | |
| inter_loc = interventions_needed[:,1:] | |
| scores = torch.zeros(len(inter_z)) | |
| batch_size = len(batch[0]) | |
| for j in range(0, len(inter_z), batch_size): | |
| ibz = inter_z[j:j+batch_size] | |
| ibl = inter_loc[j:j+batch_size].t() | |
| imask = torch.zeros((len(ibz),) + feature_shape, device=ibz.device) | |
| imask[(torch.arange(len(ibz)),) + tuple(ibl)] = 1 | |
| ibc = inter_chan[j:j+batch_size] | |
| model.ablation[layer] = ( | |
| imask.float()[:,None,:,:] * ibc[:,:,None,None]) | |
| tensor_images = model(ibz) | |
| seg = segmenter.segment_batch(tensor_images, downsample=2) | |
| mask = (seg == classnum).max(1)[0] | |
| downsampled_iseg = torch.nn.functional.adaptive_avg_pool2d( | |
| mask.float()[:,None,:,:], feature_shape)[:,0,:,:] | |
| scores[j:j+batch_size] = downsampled_iseg[ | |
| (torch.arange(len(ibz)),) + tuple(ibl)] | |
| scores = scores.view(repeats, location_count).sum(1) | |
| total_scores[1:] += scores | |
| return total_scores | |
| def count_segments(segmenter, loader, model): | |
| total_bincount = 0 | |
| data_size = 0 | |
| progress = default_progress() | |
| for i, batch in enumerate(progress(loader)): | |
| tensor_images = model(z_batch.to(device)) | |
| seg = segmenter.segment_batch(tensor_images, downsample=2) | |
| bc = (seg + index[:, None, None, None] * self.num_classes).view(-1 | |
| ).bincount(minlength=z_batch.shape[0] * self.num_classes) | |
| data_size += seg.shape[0] * seg.shape[2] * seg.shape[3] | |
| total_bincount += batch_label_counts.float().sum(0) | |
| normalized_bincount = total_bincount / data_size | |
| return normalized_bincount | |
| if __name__ == '__main__': | |
| main() | |