Spaces:
Build error
Build error
| # Instantiate the segmenter gadget. | |
| # Instantiate the GAN to optimize over | |
| # Instrument the GAN for editing and optimization. | |
| # Read quantile stats to learn 99.9th percentile for each unit, | |
| # and also the 0.01th percentile. | |
| # Read the median activation conditioned on door presence. | |
| import os, sys, numpy, torch, argparse, skimage, json, shutil | |
| from PIL import Image | |
| from torch.utils.data import TensorDataset | |
| from matplotlib.figure import Figure | |
| from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas | |
| import matplotlib.gridspec as gridspec | |
| from scipy.ndimage.morphology import binary_dilation | |
| import netdissect.zdataset | |
| import netdissect.nethook | |
| from netdissect.dissection import safe_dir_name | |
| from netdissect.progress import verbose_progress, default_progress | |
| from netdissect.progress import print_progress, desc_progress, post_progress | |
| from netdissect.easydict import EasyDict | |
| from netdissect.workerpool import WorkerPool, WorkerBase | |
| from netdissect.runningstats import RunningQuantile | |
| from netdissect.pidfile import pidfile_taken | |
| from netdissect.modelconfig import create_instrumented_model | |
| from netdissect.autoeval import autoimport_eval | |
| def main(): | |
| parser = argparse.ArgumentParser(description='ACE optimization utility', | |
| prog='python -m netdissect.aceoptimize') | |
| 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('--segmenter', type=str, default=None, | |
| help='constructor for asegmenter class') | |
| parser.add_argument('--classname', type=str, default=None, | |
| help='intervention classname') | |
| parser.add_argument('--layer', type=str, default='layer4', | |
| help='layer name') | |
| parser.add_argument('--search_size', type=int, default=10000, | |
| help='size of search for finding training locations') | |
| parser.add_argument('--train_size', type=int, default=1000, | |
| help='size of training set') | |
| parser.add_argument('--eval_size', type=int, default=200, | |
| help='size of eval set') | |
| parser.add_argument('--inference_batch_size', type=int, default=10, | |
| help='forward pass batch size') | |
| parser.add_argument('--train_batch_size', type=int, default=2, | |
| help='backprop pass batch size') | |
| parser.add_argument('--train_update_freq', type=int, default=10, | |
| help='number of batches for each training update') | |
| parser.add_argument('--train_epochs', type=int, default=10, | |
| help='number of epochs of training') | |
| parser.add_argument('--l2_lambda', type=float, default=0.005, | |
| help='l2 regularizer hyperparameter') | |
| parser.add_argument('--eval_only', action='store_true', default=False, | |
| help='reruns eval only on trained snapshots') | |
| parser.add_argument('--no-cuda', action='store_true', default=False, | |
| help='disables CUDA usage') | |
| parser.add_argument('--no-cache', action='store_true', default=False, | |
| help='disables reading of cache') | |
| parser.add_argument('--outdir', type=str, default=None, | |
| help='dissection directory') | |
| parser.add_argument('--variant', type=str, default=None, | |
| help='experiment variant') | |
| args = parser.parse_args() | |
| args.cuda = not args.no_cuda and torch.cuda.is_available() | |
| torch.backends.cudnn.benchmark = True | |
| run_command(args) | |
| def run_command(args): | |
| verbose_progress(True) | |
| progress = default_progress() | |
| classname = args.classname # 'door' | |
| layer = args.layer # 'layer4' | |
| num_eval_units = 20 | |
| assert os.path.isfile(os.path.join(args.outdir, 'dissect.json')), ( | |
| "Should be a dissection directory") | |
| if args.variant is None: | |
| args.variant = 'ace' | |
| if args.l2_lambda != 0.005: | |
| args.variant = '%s_reg%g' % (args.variant, args.l2_lambda) | |
| cachedir = os.path.join(args.outdir, safe_dir_name(layer), args.variant, | |
| classname) | |
| if pidfile_taken(os.path.join(cachedir, 'lock.pid'), True): | |
| sys.exit(0) | |
| # 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 | |
| # Default segmenter class | |
| if args.segmenter is None: | |
| args.segmenter = ("netdissect.segmenter.UnifiedParsingSegmenter(" + | |
| "segsizes=[256], segdiv='quad')") | |
| if (not args.no_cache and | |
| os.path.isfile(os.path.join(cachedir, 'snapshots', 'epoch-%d.npy' % ( | |
| args.train_epochs - 1))) and | |
| os.path.isfile(os.path.join(cachedir, 'report.json'))): | |
| print('%s already done' % cachedir) | |
| sys.exit(0) | |
| os.makedirs(cachedir, exist_ok=True) | |
| # Instantiate generator | |
| model = create_instrumented_model(args, gen=True, edit=True, | |
| layers=[args.layer]) | |
| if model is None: | |
| print('No model specified') | |
| sys.exit(1) | |
| # Instantiate segmenter | |
| segmenter = autoimport_eval(args.segmenter) | |
| labelnames, catname = segmenter.get_label_and_category_names() | |
| classnum = [i for i, (n, c) in enumerate(labelnames) if n == classname][0] | |
| num_classes = len(labelnames) | |
| with open(os.path.join(cachedir, 'labelnames.json'), 'w') as f: | |
| json.dump(labelnames, f, indent=1) | |
| # Sample sets for training. | |
| full_sample = netdissect.zdataset.z_sample_for_model(model, | |
| args.search_size, seed=10) | |
| second_sample = netdissect.zdataset.z_sample_for_model(model, | |
| args.search_size, seed=11) | |
| # Load any cached data. | |
| cache_filename = os.path.join(cachedir, 'corpus.npz') | |
| corpus = EasyDict() | |
| try: | |
| if not args.no_cache: | |
| corpus = EasyDict({k: torch.from_numpy(v) | |
| for k, v in numpy.load(cache_filename).items()}) | |
| except: | |
| pass | |
| # The steps for the computation. | |
| compute_present_locations(args, corpus, cache_filename, | |
| model, segmenter, classnum, full_sample) | |
| compute_mean_present_features(args, corpus, cache_filename, model) | |
| compute_feature_quantiles(args, corpus, cache_filename, model, full_sample) | |
| compute_candidate_locations(args, corpus, cache_filename, model, segmenter, | |
| classnum, second_sample) | |
| # visualize_training_locations(args, corpus, cachedir, model) | |
| init_ablation = initial_ablation(args, args.outdir) | |
| scores = train_ablation(args, corpus, cache_filename, | |
| model, segmenter, classnum, init_ablation) | |
| summarize_scores(args, corpus, cachedir, layer, classname, | |
| args.variant, scores) | |
| if args.variant == 'ace': | |
| add_ace_ranking_to_dissection(args.outdir, layer, classname, scores) | |
| # TODO: do some evaluation. | |
| class SaveImageWorker(WorkerBase): | |
| def work(self, data, filename): | |
| Image.fromarray(data).save(filename, optimize=True, quality=80) | |
| def plot_heatmap(output_filename, data, size=256): | |
| fig = Figure(figsize=(1, 1), dpi=size) | |
| canvas = FigureCanvas(fig) | |
| gs = gridspec.GridSpec(1, 1, left=0.0, right=1.0, bottom=0.0, top=1.0) | |
| ax = fig.add_subplot(gs[0]) | |
| ax.set_axis_off() | |
| ax.imshow(data, cmap='hot', aspect='equal', interpolation='nearest', | |
| vmin=-1, vmax=1) | |
| canvas.print_figure(output_filename, format='png') | |
| def draw_heatmap(output_filename, data, size=256): | |
| fig = Figure(figsize=(1, 1), dpi=size) | |
| canvas = FigureCanvas(fig) | |
| gs = gridspec.GridSpec(1, 1, left=0.0, right=1.0, bottom=0.0, top=1.0) | |
| ax = fig.add_subplot(gs[0]) | |
| ax.set_axis_off() | |
| ax.imshow(data, cmap='hot', aspect='equal', interpolation='nearest', | |
| vmin=-1, vmax=1) | |
| canvas.draw() # draw the canvas, cache the renderer | |
| image = numpy.fromstring(canvas.tostring_rgb(), dtype='uint8').reshape( | |
| (size, size, 3)) | |
| return image | |
| def compute_present_locations(args, corpus, cache_filename, | |
| model, segmenter, classnum, full_sample): | |
| # Phase 1. Identify a set of locations where there are doorways. | |
| # Segment the image and find featuremap pixels that maximize the number | |
| # of doorway pixels under the featuremap pixel. | |
| if all(k in corpus for k in ['present_indices', | |
| 'object_present_sample', 'object_present_location', | |
| 'object_location_popularity', 'weighted_mean_present_feature']): | |
| return | |
| progress = default_progress() | |
| feature_shape = model.feature_shape[args.layer][2:] | |
| num_locations = numpy.prod(feature_shape).item() | |
| num_units = model.feature_shape[args.layer][1] | |
| with torch.no_grad(): | |
| weighted_feature_sum = torch.zeros(num_units).cuda() | |
| object_presence_scores = [] | |
| for [zbatch] in progress( | |
| torch.utils.data.DataLoader(TensorDataset(full_sample), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Object pool"): | |
| zbatch = zbatch.cuda() | |
| tensor_image = model(zbatch) | |
| segmented_image = segmenter.segment_batch(tensor_image, | |
| downsample=2) | |
| mask = (segmented_image == classnum).max(1)[0] | |
| score = torch.nn.functional.adaptive_avg_pool2d( | |
| mask.float(), feature_shape) | |
| object_presence_scores.append(score.cpu()) | |
| feat = model.retained_layer(args.layer) | |
| weighted_feature_sum += (feat * score[:,None,:,:]).view( | |
| feat.shape[0],feat.shape[1], -1).sum(2).sum(0) | |
| object_presence_at_feature = torch.cat(object_presence_scores) | |
| object_presence_at_image, object_location_in_image = ( | |
| object_presence_at_feature.view(args.search_size, -1).max(1)) | |
| best_presence_scores, best_presence_images = torch.sort( | |
| -object_presence_at_image) | |
| all_present_indices = torch.sort( | |
| best_presence_images[:(args.train_size+args.eval_size)])[0] | |
| corpus.present_indices = all_present_indices[:args.train_size] | |
| corpus.object_present_sample = full_sample[corpus.present_indices] | |
| corpus.object_present_location = object_location_in_image[ | |
| corpus.present_indices] | |
| corpus.object_location_popularity = torch.bincount( | |
| corpus.object_present_location, | |
| minlength=num_locations) | |
| corpus.weighted_mean_present_feature = (weighted_feature_sum.cpu() / ( | |
| 1e-20 + object_presence_at_feature.view(-1).sum())) | |
| corpus.eval_present_indices = all_present_indices[-args.eval_size:] | |
| corpus.eval_present_sample = full_sample[corpus.eval_present_indices] | |
| corpus.eval_present_location = object_location_in_image[ | |
| corpus.eval_present_indices] | |
| if cache_filename: | |
| numpy.savez(cache_filename, **corpus) | |
| def compute_mean_present_features(args, corpus, cache_filename, model): | |
| # Phase 1.5. Figure mean activations for every channel where there | |
| # is a doorway. | |
| if all(k in corpus for k in ['mean_present_feature']): | |
| return | |
| progress = default_progress() | |
| with torch.no_grad(): | |
| total_present_feature = 0 | |
| for [zbatch, featloc] in progress( | |
| torch.utils.data.DataLoader(TensorDataset( | |
| corpus.object_present_sample, | |
| corpus.object_present_location), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Mean activations"): | |
| zbatch = zbatch.cuda() | |
| featloc = featloc.cuda() | |
| tensor_image = model(zbatch) | |
| feat = model.retained_layer(args.layer) | |
| flatfeat = feat.view(feat.shape[0], feat.shape[1], -1) | |
| sum_feature_at_obj = flatfeat[ | |
| torch.arange(feat.shape[0]).to(feat.device), :, featloc | |
| ].sum(0) | |
| total_present_feature = total_present_feature + sum_feature_at_obj | |
| corpus.mean_present_feature = (total_present_feature / len( | |
| corpus.object_present_sample)).cpu() | |
| if cache_filename: | |
| numpy.savez(cache_filename, **corpus) | |
| def compute_feature_quantiles(args, corpus, cache_filename, model, full_sample): | |
| # Phase 1.6. Figure the 99% and 99.9%ile of every feature. | |
| if all(k in corpus for k in ['feature_99', 'feature_999']): | |
| return | |
| progress = default_progress() | |
| with torch.no_grad(): | |
| rq = RunningQuantile(resolution=10000) # 10x what's needed. | |
| for [zbatch] in progress( | |
| torch.utils.data.DataLoader(TensorDataset(full_sample), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Calculating 0.999 quantile"): | |
| zbatch = zbatch.cuda() | |
| tensor_image = model(zbatch) | |
| feat = model.retained_layer(args.layer) | |
| rq.add(feat.permute(0, 2, 3, 1 | |
| ).contiguous().view(-1, feat.shape[1])) | |
| result = rq.quantiles([0.001, 0.01, 0.1, 0.5, 0.9, 0.99, 0.999]) | |
| corpus.feature_001 = result[:, 0].cpu() | |
| corpus.feature_01 = result[:, 1].cpu() | |
| corpus.feature_10 = result[:, 2].cpu() | |
| corpus.feature_50 = result[:, 3].cpu() | |
| corpus.feature_90 = result[:, 4].cpu() | |
| corpus.feature_99 = result[:, 5].cpu() | |
| corpus.feature_999 = result[:, 6].cpu() | |
| numpy.savez(cache_filename, **corpus) | |
| def compute_candidate_locations(args, corpus, cache_filename, model, | |
| segmenter, classnum, second_sample): | |
| # Phase 2. Identify a set of candidate locations for doorways. | |
| # Place the median doorway activation in every location of an image | |
| # and identify where it can go that doorway pixels increase. | |
| if all(k in corpus for k in ['candidate_indices', | |
| 'candidate_sample', 'candidate_score', | |
| 'candidate_location', 'object_score_at_candidate', | |
| 'candidate_location_popularity']): | |
| return | |
| progress = default_progress() | |
| feature_shape = model.feature_shape[args.layer][2:] | |
| num_locations = numpy.prod(feature_shape).item() | |
| with torch.no_grad(): | |
| # Simplify - just treat all locations as possible | |
| possible_locations = numpy.arange(num_locations) | |
| # Speed up search for locations, by weighting probed locations | |
| # according to observed distribution. | |
| location_weights = (corpus.object_location_popularity).double() | |
| location_weights += (location_weights.mean()) / 10.0 | |
| location_weights = location_weights / location_weights.sum() | |
| candidate_scores = [] | |
| object_scores = [] | |
| prng = numpy.random.RandomState(1) | |
| for [zbatch] in progress( | |
| torch.utils.data.DataLoader(TensorDataset(second_sample), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Candidate pool"): | |
| batch_scores = torch.zeros((len(zbatch),) + feature_shape).cuda() | |
| flat_batch_scores = batch_scores.view(len(zbatch), -1) | |
| zbatch = zbatch.cuda() | |
| tensor_image = model(zbatch) | |
| segmented_image = segmenter.segment_batch(tensor_image, | |
| downsample=2) | |
| mask = (segmented_image == classnum).max(1)[0] | |
| object_score = torch.nn.functional.adaptive_avg_pool2d( | |
| mask.float(), feature_shape) | |
| baseline_presence = mask.float().view(mask.shape[0], -1).sum(1) | |
| edit_mask = torch.zeros((1, 1) + feature_shape).cuda() | |
| if '_tcm' in args.variant: | |
| # variant: top-conditional-mean | |
| replace_vec = (corpus.mean_present_feature | |
| [None,:,None,None].cuda()) | |
| else: # default: weighted mean | |
| replace_vec = (corpus.weighted_mean_present_feature | |
| [None,:,None,None].cuda()) | |
| # Sample 10 random locations to examine. | |
| for loc in prng.choice(possible_locations, replace=False, | |
| p=location_weights, size=5): | |
| edit_mask.zero_() | |
| edit_mask.view(-1)[loc] = 1 | |
| model.edit_layer(args.layer, | |
| ablation=edit_mask, replacement=replace_vec) | |
| tensor_image = model(zbatch) | |
| segmented_image = segmenter.segment_batch(tensor_image, | |
| downsample=2) | |
| mask = (segmented_image == classnum).max(1)[0] | |
| modified_presence = mask.float().view( | |
| mask.shape[0], -1).sum(1) | |
| flat_batch_scores[:,loc] = ( | |
| modified_presence - baseline_presence) | |
| candidate_scores.append(batch_scores.cpu()) | |
| object_scores.append(object_score.cpu()) | |
| object_scores = torch.cat(object_scores) | |
| candidate_scores = torch.cat(candidate_scores) | |
| # Eliminate candidates where the object is present. | |
| candidate_scores = candidate_scores * (object_scores == 0).float() | |
| candidate_score_at_image, candidate_location_in_image = ( | |
| candidate_scores.view(args.search_size, -1).max(1)) | |
| best_candidate_scores, best_candidate_images = torch.sort( | |
| -candidate_score_at_image) | |
| all_candidate_indices = torch.sort( | |
| best_candidate_images[:(args.train_size+args.eval_size)])[0] | |
| corpus.candidate_indices = all_candidate_indices[:args.train_size] | |
| corpus.candidate_sample = second_sample[corpus.candidate_indices] | |
| corpus.candidate_location = candidate_location_in_image[ | |
| corpus.candidate_indices] | |
| corpus.candidate_score = candidate_score_at_image[ | |
| corpus.candidate_indices] | |
| corpus.object_score_at_candidate = object_scores.view( | |
| len(object_scores), -1)[ | |
| corpus.candidate_indices, corpus.candidate_location] | |
| corpus.candidate_location_popularity = torch.bincount( | |
| corpus.candidate_location, | |
| minlength=num_locations) | |
| corpus.eval_candidate_indices = all_candidate_indices[ | |
| -args.eval_size:] | |
| corpus.eval_candidate_sample = second_sample[ | |
| corpus.eval_candidate_indices] | |
| corpus.eval_candidate_location = candidate_location_in_image[ | |
| corpus.eval_candidate_indices] | |
| numpy.savez(cache_filename, **corpus) | |
| def visualize_training_locations(args, corpus, cachedir, model): | |
| # Phase 2.5 Create visualizations of the corpus images. | |
| progress = default_progress() | |
| feature_shape = model.feature_shape[args.layer][2:] | |
| num_locations = numpy.prod(feature_shape).item() | |
| with torch.no_grad(): | |
| imagedir = os.path.join(cachedir, 'image') | |
| os.makedirs(imagedir, exist_ok=True) | |
| image_saver = WorkerPool(SaveImageWorker) | |
| for group, group_sample, group_location, group_indices in [ | |
| ('present', | |
| corpus.object_present_sample, | |
| corpus.object_present_location, | |
| corpus.present_indices), | |
| ('candidate', | |
| corpus.candidate_sample, | |
| corpus.candidate_location, | |
| corpus.candidate_indices)]: | |
| for [zbatch, featloc, indices] in progress( | |
| torch.utils.data.DataLoader(TensorDataset( | |
| group_sample, group_location, group_indices), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Visualize %s" % group): | |
| zbatch = zbatch.cuda() | |
| tensor_image = model(zbatch) | |
| feature_mask = torch.zeros((len(zbatch), 1) + feature_shape) | |
| feature_mask.view(len(zbatch), -1).scatter_( | |
| 1, featloc[:,None], 1) | |
| feature_mask = torch.nn.functional.adaptive_max_pool2d( | |
| feature_mask.float(), tensor_image.shape[-2:]).cuda() | |
| yellow = torch.Tensor([1.0, 1.0, -1.0] | |
| )[None, :, None, None].cuda() | |
| tensor_image = tensor_image * (1 - 0.5 * feature_mask) + ( | |
| 0.5 * feature_mask * yellow) | |
| byte_image = (((tensor_image+1)/2)*255).clamp(0, 255).byte() | |
| numpy_image = byte_image.permute(0, 2, 3, 1).cpu().numpy() | |
| for i, index in enumerate(indices): | |
| image_saver.add(numpy_image[i], os.path.join(imagedir, | |
| '%s_%d.jpg' % (group, index))) | |
| image_saver.join() | |
| def scale_summary(scale, lownums, highnums): | |
| value, order = (-(scale.detach())).cpu().sort(0) | |
| lowsum = ' '.join('%d: %.3g' % (o.item(), -v.item()) | |
| for v, o in zip(value[:lownums], order[:lownums])) | |
| highsum = ' '.join('%d: %.3g' % (o.item(), -v.item()) | |
| for v, o in zip(value[-highnums:], order[-highnums:])) | |
| return lowsum + ' ... ' + highsum | |
| # Phase 3. Given those two sets, now optimize a such that: | |
| # Door pred lost if we take 0 * a at a candidate (1) | |
| # Door pred gained If we take 99.9th activation * a at a candiate (1) | |
| # | |
| # ADE_au = E | on - E | off) | |
| # = cand-frac E_cand | on + nocand-frac E_cand | on | |
| # - door-frac E_door | off + nodoor-frac E_nodoor | off | |
| # approx = cand-frac E_cand | on - door-frac E_door | off + K | |
| # Each batch has both types, and minimizes | |
| # door-frac sum(s_c) when pixel off - cand-frac sum(s_c) when pixel on | |
| def initial_ablation(args, dissectdir): | |
| # Load initialization from dissection, based on iou scores. | |
| with open(os.path.join(dissectdir, 'dissect.json')) as f: | |
| dissection = EasyDict(json.load(f)) | |
| lrec = [l for l in dissection.layers if l.layer == args.layer][0] | |
| rrec = [r for r in lrec.rankings if r.name == '%s-iou' % args.classname | |
| ][0] | |
| init_scores = -torch.tensor(rrec.score) | |
| return init_scores / init_scores.max() | |
| def ace_loss(segmenter, classnum, model, layer, high_replacement, ablation, | |
| pbatch, ploc, cbatch, cloc, run_backward=False, | |
| discrete_pixels=False, | |
| discrete_units=False, | |
| mixed_units=False, | |
| ablation_only=False, | |
| fullimage_measurement=False, | |
| fullimage_ablation=False, | |
| ): | |
| feature_shape = model.feature_shape[layer][2:] | |
| if discrete_units: # discretize ablation to the top N units | |
| assert discrete_units > 0 | |
| d = torch.zeros_like(ablation) | |
| top_units = torch.topk(ablation.view(-1), discrete_units)[1] | |
| if mixed_units: | |
| d.view(-1)[top_units] = ablation.view(-1)[top_units] | |
| else: | |
| d.view(-1)[top_units] = 1 | |
| ablation = d | |
| # First, ablate a sample of locations with positive presence | |
| # and see how much the presence is reduced. | |
| p_mask = torch.zeros((len(pbatch), 1) + feature_shape) | |
| if fullimage_ablation: | |
| p_mask[...] = 1 | |
| else: | |
| p_mask.view(len(pbatch), -1).scatter_(1, ploc[:,None], 1) | |
| p_mask = p_mask.cuda() | |
| a_p_mask = (ablation * p_mask) | |
| model.edit_layer(layer, ablation=a_p_mask, replacement=None) | |
| tensor_images = model(pbatch.cuda()) | |
| assert model._ablation[layer] is a_p_mask | |
| erase_effect, erased_mask = segmenter.predict_single_class( | |
| tensor_images, classnum, downsample=2) | |
| if discrete_pixels: # pixel loss: use mask instead of pred | |
| erase_effect = erased_mask.float() | |
| erase_downsampled = torch.nn.functional.adaptive_avg_pool2d( | |
| erase_effect[:,None,:,:], feature_shape)[:,0,:,:] | |
| if fullimage_measurement: | |
| erase_loss = erase_downsampled.sum() | |
| else: | |
| erase_at_loc = erase_downsampled.view(len(erase_downsampled), -1 | |
| )[torch.arange(len(erase_downsampled)), ploc] | |
| erase_loss = erase_at_loc.sum() | |
| if run_backward: | |
| erase_loss.backward() | |
| if ablation_only: | |
| return erase_loss | |
| # Second, activate a sample of locations that are candidates for | |
| # insertion and see how much the presence is increased. | |
| c_mask = torch.zeros((len(cbatch), 1) + feature_shape) | |
| c_mask.view(len(cbatch), -1).scatter_(1, cloc[:,None], 1) | |
| c_mask = c_mask.cuda() | |
| a_c_mask = (ablation * c_mask) | |
| model.edit_layer(layer, ablation=a_c_mask, replacement=high_replacement) | |
| tensor_images = model(cbatch.cuda()) | |
| assert model._ablation[layer] is a_c_mask | |
| add_effect, added_mask = segmenter.predict_single_class( | |
| tensor_images, classnum, downsample=2) | |
| if discrete_pixels: # pixel loss: use mask instead of pred | |
| add_effect = added_mask.float() | |
| add_effect = -add_effect | |
| add_downsampled = torch.nn.functional.adaptive_avg_pool2d( | |
| add_effect[:,None,:,:], feature_shape)[:,0,:,:] | |
| if fullimage_measurement: | |
| add_loss = add_downsampled.mean() | |
| else: | |
| add_at_loc = add_downsampled.view(len(add_downsampled), -1 | |
| )[torch.arange(len(add_downsampled)), ploc] | |
| add_loss = add_at_loc.sum() | |
| if run_backward: | |
| add_loss.backward() | |
| return erase_loss + add_loss | |
| 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() | |
| if '_h99' in args.variant: | |
| high_replacement = corpus.feature_99[None,:,None,None].cuda() | |
| elif '_tcm' in args.variant: | |
| # variant: top-conditional-mean | |
| high_replacement = ( | |
| corpus.mean_present_feature[None,:,None,None].cuda()) | |
| else: # default: weighted mean | |
| high_replacement = ( | |
| corpus.weighted_mean_present_feature[None,:,None,None].cuda()) | |
| fullimage_measurement = False | |
| ablation_only = False | |
| fullimage_ablation = False | |
| if '_fim' in args.variant: | |
| fullimage_measurement = True | |
| elif '_fia' in args.variant: | |
| fullimage_measurement = True | |
| ablation_only = True | |
| fullimage_ablation = True | |
| 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 | |
| epoch = 0 | |
| 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 | |
| if args.eval_only: | |
| # For eval_only, just load each snapshot and re-run validation eval | |
| # pass on each one. | |
| for epoch in range(-1, args.train_epochs): | |
| snapfile = os.path.join(snapdir, 'epoch-%d.pth' % epoch) | |
| if not os.path.exists(snapfile): | |
| data = {} | |
| if epoch >= 0: | |
| print('No epoch %d' % epoch) | |
| continue | |
| else: | |
| data = torch.load(snapfile) | |
| with torch.no_grad(): | |
| ablation[...] = data['ablation'].to(ablation.device) | |
| optimizer.load_state_dict(data['optimizer']) | |
| avg_loss, regularizer, new_extra = eval_loss_and_reg() | |
| # Keep old values, and update any new ones. | |
| extra = {k: v for k, v in data.items() | |
| if k not in ['ablation', 'optimizer', 'avg_loss']} | |
| extra.update(new_extra) | |
| torch.save(dict(ablation=ablation, optimizer=optimizer.state_dict(), | |
| avg_loss=avg_loss, **extra), | |
| os.path.join(snapdir, 'epoch-%d.pth' % epoch)) | |
| # Return loaded ablation. | |
| return ablation.view(-1).detach().cpu().numpy() | |
| 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 | |
| if start_epoch < args.train_epochs: | |
| epoch = start_epoch - 1 | |
| avg_loss, regularizer, extra = eval_loss_and_reg() | |
| if epoch == -1: | |
| torch.save(dict(ablation=ablation, optimizer=optimizer.state_dict(), | |
| avg_loss=avg_loss, **extra), | |
| os.path.join(snapdir, 'epoch-%d.pth' % epoch)) | |
| 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, | |
| ablation_only=ablation_only, | |
| fullimage_measurement=fullimage_measurement) | |
| 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 | |
| avg_loss, regularizer, extra = eval_loss_and_reg() | |
| torch.save(dict(ablation=ablation, optimizer=optimizer.state_dict(), | |
| avg_loss=avg_loss, **extra), | |
| 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 tensor_to_numpy_image_batch(tensor_image): | |
| byte_image = (((tensor_image+1)/2)*255).clamp(0, 255).byte() | |
| numpy_image = byte_image.permute(0, 2, 3, 1).cpu().numpy() | |
| return numpy_image | |
| # Phase 4: evaluation of intervention | |
| def evaluate_ablation(args, model, segmenter, eval_sample, classnum, layer, | |
| ordering): | |
| total_bincount = 0 | |
| data_size = 0 | |
| progress = default_progress() | |
| for l in model.ablation: | |
| model.ablation[l] = None | |
| feature_units = model.feature_shape[args.layer][1] | |
| feature_shape = model.feature_shape[args.layer][2:] | |
| repeats = len(ordering) | |
| total_scores = torch.zeros(repeats + 1) | |
| for i, batch in enumerate(progress(torch.utils.data.DataLoader( | |
| TensorDataset(eval_sample), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Evaluate interventions")): | |
| tensor_image = model(zbatch) | |
| segmented_image = segmenter.segment_batch(tensor_image, | |
| downsample=2) | |
| mask = (segmented_image == 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.edit_layer(args.layer, ablation=( | |
| imask.float()[:,None,:,:] * ibc[:,:,None,None])) | |
| _, seg, _, _, _ = ( | |
| recovery.recover_im_seg_bc_and_features( | |
| [ibz], model)) | |
| 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 evaluate_interventions(args, model, segmenter, eval_sample, | |
| classnum, layer, units): | |
| total_bincount = 0 | |
| data_size = 0 | |
| progress = default_progress() | |
| for l in model.ablation: | |
| model.ablation[l] = None | |
| feature_units = model.feature_shape[args.layer][1] | |
| feature_shape = model.feature_shape[args.layer][2:] | |
| repeats = len(ordering) | |
| total_scores = torch.zeros(repeats + 1) | |
| for i, batch in enumerate(progress(torch.utils.data.DataLoader( | |
| TensorDataset(eval_sample), | |
| batch_size=args.inference_batch_size, num_workers=10, | |
| pin_memory=True), | |
| desc="Evaluate interventions")): | |
| tensor_image = model(zbatch) | |
| segmented_image = segmenter.segment_batch(tensor_image, | |
| downsample=2) | |
| mask = (segmented_image == 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[args.layer] = ( | |
| imask.float()[:,None,:,:] * ibc[:,:,None,None]) | |
| _, seg, _, _, _ = ( | |
| recovery.recover_im_seg_bc_and_features( | |
| [ibz], model)) | |
| 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 add_ace_ranking_to_dissection(outdir, layer, classname, total_scores): | |
| source_filename = os.path.join(outdir, 'dissect.json') | |
| source_filename_bak = os.path.join(outdir, 'dissect.json.bak') | |
| # Back up the dissection (if not already backed up) before modifying | |
| if not os.path.exists(source_filename_bak): | |
| shutil.copy(source_filename, source_filename_bak) | |
| with open(source_filename) as f: | |
| dissection = EasyDict(json.load(f)) | |
| ranking_name = '%s-ace' % classname | |
| # Remove any old ace ranking with the same name | |
| lrec = [l for l in dissection.layers if l.layer == layer][0] | |
| lrec.rankings = [r for r in lrec.rankings if r.name != ranking_name] | |
| # Now convert ace scores to rankings | |
| new_rankings = [dict( | |
| name=ranking_name, | |
| score=(-total_scores).flatten().tolist(), | |
| metric='ace')] | |
| # Prepend to list. | |
| lrec.rankings[2:2] = new_rankings | |
| # Replace the old dissect.json in-place | |
| with open(source_filename, 'w') as f: | |
| json.dump(dissection, f, indent=1) | |
| def summarize_scores(args, corpus, cachedir, layer, classname, variant, scores): | |
| target_filename = os.path.join(cachedir, 'summary.json') | |
| ranking_name = '%s-%s' % (classname, variant) | |
| # Now convert ace scores to rankings | |
| new_rankings = [dict( | |
| name=ranking_name, | |
| score=(-scores).flatten().tolist(), | |
| metric=variant)] | |
| result = dict(layers=[dict(layer=layer, rankings=new_rankings)]) | |
| # Replace the old dissect.json in-place | |
| with open(target_filename, 'w') as f: | |
| json.dump(result, f, indent=1) | |
| if __name__ == '__main__': | |
| main() | |