netdissect/aceoptimize.py

# 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()