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 verbose_progress, print_progress | |
| from netdissect import InstrumentedModel, BrodenDataset, dissect | |
| from netdissect import MultiSegmentDataset, GeneratorSegRunner | |
| from netdissect import ImageOnlySegRunner | |
| from netdissect.parallelfolder import ParallelImageFolders | |
| from netdissect.zdataset import z_dataset_for_model | |
| from netdissect.autoeval import autoimport_eval | |
| from netdissect.modelconfig import create_instrumented_model | |
| from netdissect.pidfile import exit_if_job_done, mark_job_done | |
| help_epilog = '''\ | |
| Example: to dissect three layers of the pretrained alexnet in torchvision: | |
| python -m netdissect \\ | |
| --model "torchvision.models.alexnet(pretrained=True)" \\ | |
| --layers features.6:conv3 features.8:conv4 features.10:conv5 \\ | |
| --imgsize 227 \\ | |
| --outdir dissect/alexnet-imagenet | |
| To dissect a progressive GAN model: | |
| python -m netdissect \\ | |
| --model "proggan.from_pth_file('model/churchoutdoor.pth')" \\ | |
| --gan | |
| ''' | |
| 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='quad')") | |
| # 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) | |
| segrunner = GeneratorSegRunner(autoimport_eval(args.segmenter)) | |
| # 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) | |
| class AddPerturbation(object): | |
| def __init__(self, perturbation): | |
| self.perturbation = perturbation | |
| def __call__(self, pic): | |
| if self.perturbation is None: | |
| return pic | |
| # Convert to a numpy float32 array | |
| npyimg = numpy.array(pic, numpy.uint8, copy=False | |
| ).astype(numpy.float32) | |
| # Center the perturbation | |
| oy, ox = ((self.perturbation.shape[d] - npyimg.shape[d]) // 2 | |
| for d in [0, 1]) | |
| npyimg += self.perturbation[ | |
| oy:oy+npyimg.shape[0], ox:ox+npyimg.shape[1]] | |
| # Pytorch conventions: as a float it should be [0..1] | |
| npyimg.clip(0, 255, npyimg) | |
| return npyimg / 255.0 | |
| def test_dissection(): | |
| verbose_progress(True) | |
| from torchvision.models import alexnet | |
| from torchvision import transforms | |
| model = InstrumentedModel(alexnet(pretrained=True)) | |
| model.eval() | |
| # Load an alexnet | |
| model.retain_layers([ | |
| ('features.0', 'conv1'), | |
| ('features.3', 'conv2'), | |
| ('features.6', 'conv3'), | |
| ('features.8', 'conv4'), | |
| ('features.10', 'conv5') ]) | |
| # load broden dataset | |
| bds = BrodenDataset('dataset/broden', | |
| transform=transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(IMAGE_MEAN, IMAGE_STDEV)]), | |
| size=100) | |
| # run dissect | |
| dissect('dissect/test', model, bds, | |
| examples_per_unit=10) | |
| def try_to_load_images(directory, imgsize, perturbation, size): | |
| # Load plain image dataset | |
| # TODO: allow other normalizations. | |
| return ParallelImageFolders( | |
| [directory], | |
| transform=transforms.Compose([ | |
| transforms.Resize(imgsize), | |
| AddPerturbation(perturbation), | |
| transforms.ToTensor(), | |
| transforms.Normalize(IMAGE_MEAN, IMAGE_STDEV)]), | |
| size=size) | |
| def try_to_load_broden(directory, imgsize, broden_version, perturbation, | |
| download, size): | |
| # Load broden dataset | |
| ds_resolution = (224 if max(imgsize) <= 224 else | |
| 227 if max(imgsize) <= 227 else 384) | |
| if not os.path.isfile(os.path.join(directory, | |
| 'broden%d_%d' % (broden_version, ds_resolution), 'index.csv')): | |
| return None | |
| return BrodenDataset(directory, | |
| resolution=ds_resolution, | |
| download=download, | |
| broden_version=broden_version, | |
| transform=transforms.Compose([ | |
| transforms.Resize(imgsize), | |
| AddPerturbation(perturbation), | |
| transforms.ToTensor(), | |
| transforms.Normalize(IMAGE_MEAN, IMAGE_STDEV)]), | |
| size=size) | |
| def try_to_load_multiseg(directory, imgsize, perturbation, size): | |
| if not os.path.isfile(os.path.join(directory, 'labelnames.json')): | |
| return None | |
| minsize = min(imgsize) if hasattr(imgsize, '__iter__') else imgsize | |
| return MultiSegmentDataset(directory, | |
| transform=(transforms.Compose([ | |
| transforms.Resize(minsize), | |
| transforms.CenterCrop(imgsize), | |
| AddPerturbation(perturbation), | |
| transforms.ToTensor(), | |
| transforms.Normalize(IMAGE_MEAN, IMAGE_STDEV)]), | |
| transforms.Compose([ | |
| transforms.Resize(minsize, interpolation=PIL.Image.NEAREST), | |
| transforms.CenterCrop(imgsize)])), | |
| size=size) | |
| def add_scale_offset_info(model, layer_names): | |
| ''' | |
| Creates a 'scale_offset' property on the model which guesses | |
| how to offset the featuremap, in cases where the convolutional | |
| padding does not exacly correspond to keeping featuremap pixels | |
| centered on the downsampled regions of the input. This mainly | |
| shows up in AlexNet: ResNet and VGG pad convolutions to keep | |
| them centered and do not need this. | |
| ''' | |
| model.scale_offset = {} | |
| seen = set() | |
| sequence = [] | |
| aka_map = {} | |
| for name in layer_names: | |
| aka = name | |
| if not isinstance(aka, str): | |
| name, aka = name | |
| aka_map[name] = aka | |
| for name, layer in model.named_modules(): | |
| sequence.append(layer) | |
| if name in aka_map: | |
| seen.add(name) | |
| aka = aka_map[name] | |
| model.scale_offset[aka] = sequence_scale_offset(sequence) | |
| for name in aka_map: | |
| assert name in seen, ('Layer %s not found' % name) | |
| def dilation_scale_offset(dilations): | |
| '''Composes a list of (k, s, p) into a single total scale and offset.''' | |
| if len(dilations) == 0: | |
| return (1, 0) | |
| scale, offset = dilation_scale_offset(dilations[1:]) | |
| kernel, stride, padding = dilations[0] | |
| scale *= stride | |
| offset *= stride | |
| offset += (kernel - 1) / 2.0 - padding | |
| return scale, offset | |
| def dilations(modulelist): | |
| '''Converts a list of modules to (kernel_size, stride, padding)''' | |
| result = [] | |
| for module in modulelist: | |
| settings = tuple(getattr(module, n, d) | |
| for n, d in (('kernel_size', 1), ('stride', 1), ('padding', 0))) | |
| settings = (((s, s) if not isinstance(s, tuple) else s) | |
| for s in settings) | |
| if settings != ((1, 1), (1, 1), (0, 0)): | |
| result.append(zip(*settings)) | |
| return zip(*result) | |
| def sequence_scale_offset(modulelist): | |
| '''Returns (yscale, yoffset), (xscale, xoffset) given a list of modules''' | |
| return tuple(dilation_scale_offset(d) for d in dilations(modulelist)) | |
| def strfloat(s): | |
| try: | |
| return float(s) | |
| except: | |
| return s | |
| class FloatRange(object): | |
| def __init__(self, start, end): | |
| self.start = start | |
| self.end = end | |
| def __eq__(self, other): | |
| return isinstance(other, float) and self.start <= other <= self.end | |
| def __repr__(self): | |
| return '[%g-%g]' % (self.start, self.end) | |
| # Many models use this normalization. | |
| IMAGE_MEAN = [0.485, 0.456, 0.406] | |
| IMAGE_STDEV = [0.229, 0.224, 0.225] | |
| if __name__ == '__main__': | |
| main() | |