Spaces:
Build error
Build error
| ''' | |
| A simple tool to generate sample of output of a GAN, | |
| subject to filtering, sorting, or intervention. | |
| ''' | |
| import torch, numpy, os, argparse, sys, shutil, errno, numbers | |
| 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 retain_layers | |
| from netdissect.runningstats import RunningTopK | |
| 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('--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 | |
| retain_layers(model, [args.layer]) | |
| model.cuda() | |
| 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_all_highest_znums( | |
| model, z_universe, args.size, seed=args.seed) | |
| save_chosen_unit_images(args.outdir, model, z_universe, indexes, | |
| lightbox=True) | |
| def get_all_highest_znums(model, z_universe, size, | |
| batch_size=10, 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) | |
| rtk = RunningTopK(k=size) | |
| 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.view( | |
| feature.shape[0], num_units, -1).max(2)[0] | |
| rtk.add(max_feature) | |
| td, ti = rtk.result() | |
| highest = ti.sort(1)[0] | |
| return highest | |
| def save_chosen_unit_images(dirname, model, z_universe, indices, | |
| shared_dir="shared_images", | |
| unitdir_template="unit_{}", | |
| name_template="image_{}.jpg", | |
| lightbox=False, batch_size=50, seed=1): | |
| all_indices = torch.unique(indices.view(-1), sorted=True) | |
| z_sample = z_universe[all_indices] | |
| progress = default_progress() | |
| sdir = os.path.join(dirname, shared_dir) | |
| created_hashdirs = set() | |
| for index in range(len(z_universe)): | |
| hd = hashdir(index) | |
| if hd not in created_hashdirs: | |
| created_hashdirs.add(hd) | |
| os.makedirs(os.path.join(sdir, hd), 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) | |
| 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 = all_indices[i + start_index].item() | |
| filename = os.path.join(sdir, hashdir(index), | |
| name_template.format(index)) | |
| saver.add(im[i].numpy(), filename) | |
| saver.join() | |
| linker = WorkerPool(MakeLinkWorker) | |
| for u in progress(range(len(indices)), desc='Making links'): | |
| udir = os.path.join(dirname, unitdir_template.format(u)) | |
| os.makedirs(udir, exist_ok=True) | |
| for r in range(indices.shape[1]): | |
| index = indices[u,r].item() | |
| fn = name_template.format(index) | |
| # sourcename = os.path.join('..', shared_dir, fn) | |
| sourcename = os.path.join(sdir, hashdir(index), fn) | |
| targname = os.path.join(udir, fn) | |
| linker.add(sourcename, targname) | |
| if lightbox: | |
| copy_lightbox_to(udir) | |
| linker.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')) | |
| def hashdir(index): | |
| # To keep the number of files the shared directory lower, split it | |
| # into 100 subdirectories named as follows. | |
| return '%02d' % (index % 100) | |
| class SaveImageWorker(WorkerBase): | |
| # Saving images can be sped up by sending jpeg encoding and | |
| # file-writing work to a pool. | |
| def work(self, data, filename): | |
| Image.fromarray(data).save(filename, optimize=True, quality=100) | |
| class MakeLinkWorker(WorkerBase): | |
| # Creating symbolic links is a bit slow and can be done faster | |
| # in parallel rather than waiting for each to be created. | |
| def work(self, sourcename, targname): | |
| try: | |
| os.link(sourcename, targname) | |
| except OSError as e: | |
| if e.errno == errno.EEXIST: | |
| os.remove(targname) | |
| os.link(sourcename, targname) | |
| else: | |
| raise | |
| class MakeSyminkWorker(WorkerBase): | |
| # Creating symbolic links is a bit slow and can be done faster | |
| # in parallel rather than waiting for each to be created. | |
| def work(self, sourcename, targname): | |
| try: | |
| os.symlink(sourcename, targname) | |
| except OSError as e: | |
| if e.errno == errno.EEXIST: | |
| os.remove(targname) | |
| os.symlink(sourcename, targname) | |
| else: | |
| raise | |
| if __name__ == '__main__': | |
| main() | |