''' A simple tool to generate sample of output of a GAN, subject to filtering, sorting, or intervention. ''' import torch, numpy, os, argparse, numbers, sys, shutil from PIL import Image from torch.utils.data import TensorDataset from netdissect.zdataset import standard_z_sample from netdissect.progress import default_progress, verbose_progress from netdissect.autoeval import autoimport_eval from netdissect.workerpool import WorkerBase, WorkerPool from netdissect.nethook import edit_layers, retain_layers 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 get_highest_znums(model, z_universe, max_units, size, batch_size=100, seed=1): # The model should have been instrumented already retained_items = list(model.retained.items()) assert len(retained_items) == 1 layer = retained_items[0][0] # By default, a 10% sample progress = default_progress() num_units = None with torch.no_grad(): # Pass 1: collect max activation stats z_loader = torch.utils.data.DataLoader(TensorDataset(z_universe), batch_size=batch_size, num_workers=2, pin_memory=True) scores = [] for [z] in progress(z_loader, desc='Finding max activations'): z = z.cuda() model(z) feature = model.retained[layer] num_units = feature.shape[1] max_feature = feature[:, max_units, ...].view( feature.shape[0], len(max_units), -1).max(2)[0] total_feature = max_feature.sum(1) scores.append(total_feature.cpu()) scores = torch.cat(scores, 0) highest = (-scores).sort(0)[1][:size].sort(0)[0] return highest def save_znum_images(dirname, model, z_sample, indexes, layer, ablated_units, name_template="image_{}.png", lightbox=False, batch_size=100, seed=1): progress = default_progress() os.makedirs(dirname, exist_ok=True) with torch.no_grad(): # Pass 2: now generate images z_loader = torch.utils.data.DataLoader(TensorDataset(z_sample), batch_size=batch_size, num_workers=2, pin_memory=True) saver = WorkerPool(SaveImageWorker) if ablated_units is not None: dims = max(2, max(ablated_units) + 1) # >=2 to avoid broadcast mask = torch.zeros(dims) mask[ablated_units] = 1 model.ablation[layer] = mask[None,:,None,None].cuda() for batch_num, [z] in enumerate(progress(z_loader, desc='Saving images')): z = z.cuda() start_index = batch_num * batch_size im = ((model(z) + 1) / 2 * 255).clamp(0, 255).byte().permute( 0, 2, 3, 1).cpu() for i in range(len(im)): index = i + start_index if indexes is not None: index = indexes[index].item() filename = os.path.join(dirname, name_template.format(index)) saver.add(im[i].numpy(), filename) saver.join() def copy_lightbox_to(dirname): srcdir = os.path.realpath( os.path.join(os.getcwd(), os.path.dirname(__file__))) shutil.copy(os.path.join(srcdir, 'lightbox.html'), os.path.join(dirname, '+lightbox.html')) class SaveImageWorker(WorkerBase): def work(self, data, filename): Image.fromarray(data).save(filename, optimize=True, quality=100) if __name__ == '__main__': main()