ClothingGAN / netdissect /aceoptimize.py
mfrashad's picture
Init code
97069e1
# 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()