Spaces:
Sleeping
Sleeping
''' | |
To run dissection: | |
1. Load up the convolutional model you wish to dissect, and wrap it in | |
an InstrumentedModel; then call imodel.retain_layers([layernames,..]) | |
to instrument the layers of interest. | |
2. Load the segmentation dataset using the BrodenDataset class; | |
use the transform_image argument to normalize images to be | |
suitable for the model, or the size argument to truncate the dataset. | |
3. Choose a directory in which to write the output, and call | |
dissect(outdir, model, dataset). | |
Example: | |
from dissect import InstrumentedModel, dissect | |
from broden import BrodenDataset | |
model = InstrumentedModel(load_my_model()) | |
model.eval() | |
model.cuda() | |
model.retain_layers(['conv1', 'conv2', 'conv3', 'conv4', 'conv5']) | |
bds = BrodenDataset('dataset/broden1_227', | |
transform_image=transforms.Compose([ | |
transforms.ToTensor(), | |
transforms.Normalize(IMAGE_MEAN, IMAGE_STDEV)]), | |
size=1000) | |
dissect('result/dissect', model, bds, | |
examples_per_unit=10) | |
''' | |
import torch, numpy, os, re, json, shutil, types, tempfile, torchvision | |
# import warnings | |
# warnings.simplefilter('error', UserWarning) | |
from PIL import Image | |
from xml.etree import ElementTree as et | |
from collections import OrderedDict, defaultdict | |
from .progress import verbose_progress, default_progress, print_progress | |
from .progress import desc_progress | |
from .runningstats import RunningQuantile, RunningTopK | |
from .runningstats import RunningCrossCovariance, RunningConditionalQuantile | |
from .sampler import FixedSubsetSampler | |
from .actviz import activation_visualization | |
from .segviz import segment_visualization, high_contrast | |
from .workerpool import WorkerBase, WorkerPool | |
from .segmenter import UnifiedParsingSegmenter | |
def dissect(outdir, model, dataset, | |
segrunner=None, | |
train_dataset=None, | |
model_segmenter=None, | |
quantile_threshold=0.005, | |
iou_threshold=0.05, | |
iqr_threshold=0.01, | |
examples_per_unit=100, | |
batch_size=100, | |
num_workers=24, | |
seg_batch_size=5, | |
make_images=True, | |
make_labels=True, | |
make_maxiou=False, | |
make_covariance=False, | |
make_report=True, | |
make_row_images=True, | |
make_single_images=False, | |
rank_all_labels=False, | |
netname=None, | |
meta=None, | |
merge=None, | |
settings=None, | |
): | |
''' | |
Runs net dissection in-memory, using pytorch, and saves visualizations | |
and metadata into outdir. | |
''' | |
assert not model.training, 'Run model.eval() before dissection' | |
if netname is None: | |
netname = type(model).__name__ | |
if segrunner is None: | |
segrunner = ClassifierSegRunner(dataset) | |
if train_dataset is None: | |
train_dataset = dataset | |
make_iqr = (quantile_threshold == 'iqr') | |
with torch.no_grad(): | |
device = next(model.parameters()).device | |
levels = None | |
labelnames, catnames = None, None | |
maxioudata, iqrdata = None, None | |
labeldata = None | |
iqrdata, cov = None, None | |
labelnames, catnames = segrunner.get_label_and_category_names() | |
label_category = [catnames.index(c) if c in catnames else 0 | |
for l, c in labelnames] | |
# First, always collect qunatiles and topk information. | |
segloader = torch.utils.data.DataLoader(dataset, | |
batch_size=batch_size, num_workers=num_workers, | |
pin_memory=(device.type == 'cuda')) | |
quantiles, topk = collect_quantiles_and_topk(outdir, model, | |
segloader, segrunner, k=examples_per_unit) | |
# Thresholds can be automatically chosen by maximizing iqr | |
if make_iqr: | |
# Get thresholds based on an IQR optimization | |
segloader = torch.utils.data.DataLoader(train_dataset, | |
batch_size=1, num_workers=num_workers, | |
pin_memory=(device.type == 'cuda')) | |
iqrdata = collect_iqr(outdir, model, segloader, segrunner) | |
max_iqr, full_iqr_levels = iqrdata[:2] | |
max_iqr_agreement = iqrdata[4] | |
# qualified_iqr[max_iqr_quantile[layer] > 0.5] = 0 | |
levels = {layer: full_iqr_levels[layer][ | |
max_iqr[layer].max(0)[1], | |
torch.arange(max_iqr[layer].shape[1])].to(device) | |
for layer in full_iqr_levels} | |
else: | |
levels = {k: qc.quantiles([1.0 - quantile_threshold])[:,0] | |
for k, qc in quantiles.items()} | |
quantiledata = (topk, quantiles, levels, quantile_threshold) | |
if make_images: | |
segloader = torch.utils.data.DataLoader(dataset, | |
batch_size=batch_size, num_workers=num_workers, | |
pin_memory=(device.type == 'cuda')) | |
generate_images(outdir, model, dataset, topk, levels, segrunner, | |
row_length=examples_per_unit, batch_size=seg_batch_size, | |
row_images=make_row_images, | |
single_images=make_single_images, | |
num_workers=num_workers) | |
if make_maxiou: | |
assert train_dataset, "Need training dataset for maxiou." | |
segloader = torch.utils.data.DataLoader(train_dataset, | |
batch_size=1, num_workers=num_workers, | |
pin_memory=(device.type == 'cuda')) | |
maxioudata = collect_maxiou(outdir, model, segloader, | |
segrunner) | |
if make_labels: | |
segloader = torch.utils.data.DataLoader(dataset, | |
batch_size=1, num_workers=num_workers, | |
pin_memory=(device.type == 'cuda')) | |
iou_scores, iqr_scores, tcs, lcs, ccs, ics = ( | |
collect_bincounts(outdir, model, segloader, | |
levels, segrunner)) | |
labeldata = (iou_scores, iqr_scores, lcs, ccs, ics, iou_threshold, | |
iqr_threshold) | |
if make_covariance: | |
segloader = torch.utils.data.DataLoader(dataset, | |
batch_size=seg_batch_size, | |
num_workers=num_workers, | |
pin_memory=(device.type == 'cuda')) | |
cov = collect_covariance(outdir, model, segloader, segrunner) | |
if make_report: | |
generate_report(outdir, | |
quantiledata=quantiledata, | |
labelnames=labelnames, | |
catnames=catnames, | |
labeldata=labeldata, | |
maxioudata=maxioudata, | |
iqrdata=iqrdata, | |
covariancedata=cov, | |
rank_all_labels=rank_all_labels, | |
netname=netname, | |
meta=meta, | |
mergedata=merge, | |
settings=settings) | |
return quantiledata, labeldata | |
def generate_report(outdir, quantiledata, labelnames=None, catnames=None, | |
labeldata=None, maxioudata=None, iqrdata=None, covariancedata=None, | |
rank_all_labels=False, netname='Model', meta=None, settings=None, | |
mergedata=None): | |
''' | |
Creates dissection.json reports and summary bargraph.svg files in the | |
specified output directory, and copies a dissection.html interface | |
to go along with it. | |
''' | |
all_layers = [] | |
# Current source code directory, for html to copy. | |
srcdir = os.path.realpath( | |
os.path.join(os.getcwd(), os.path.dirname(__file__))) | |
# Unpack arguments | |
topk, quantiles, levels, quantile_threshold = quantiledata | |
top_record = dict( | |
netname=netname, | |
meta=meta, | |
default_ranking='unit', | |
quantile_threshold=quantile_threshold) | |
if settings is not None: | |
top_record['settings'] = settings | |
if labeldata is not None: | |
iou_scores, iqr_scores, lcs, ccs, ics, iou_threshold, iqr_threshold = ( | |
labeldata) | |
catorder = {'object': -7, 'scene': -6, 'part': -5, | |
'piece': -4, | |
'material': -3, 'texture': -2, 'color': -1} | |
for i, cat in enumerate(c for c in catnames if c not in catorder): | |
catorder[cat] = i | |
catnumber = {n: i for i, n in enumerate(catnames)} | |
catnumber['-'] = 0 | |
top_record['default_ranking'] = 'label' | |
top_record['iou_threshold'] = iou_threshold | |
top_record['iqr_threshold'] = iqr_threshold | |
labelnumber = dict((name[0], num) | |
for num, name in enumerate(labelnames)) | |
# Make a segmentation color dictionary | |
segcolors = {} | |
for i, name in enumerate(labelnames): | |
key = ','.join(str(s) for s in high_contrast[i % len(high_contrast)]) | |
if key in segcolors: | |
segcolors[key] += '/' + name[0] | |
else: | |
segcolors[key] = name[0] | |
top_record['segcolors'] = segcolors | |
for layer in topk.keys(): | |
units, rankings = [], [] | |
record = dict(layer=layer, units=units, rankings=rankings) | |
# For every unit, we always have basic visualization information. | |
topa, topi = topk[layer].result() | |
lev = levels[layer] | |
for u in range(len(topa)): | |
units.append(dict( | |
unit=u, | |
interp=True, | |
level=lev[u].item(), | |
top=[dict(imgnum=i.item(), maxact=a.item()) | |
for i, a in zip(topi[u], topa[u])], | |
)) | |
rankings.append(dict(name="unit", score=list([ | |
u for u in range(len(topa))]))) | |
# TODO: consider including stats and ranking based on quantiles, | |
# variance, connectedness here. | |
# if we have labeldata, then every unit also gets a bunch of other info | |
if labeldata is not None: | |
lscore, qscore, cc, ic = [dat[layer] | |
for dat in [iou_scores, iqr_scores, ccs, ics]] | |
if iqrdata is not None: | |
# If we have IQR thresholds, assign labels based on that | |
max_iqr, max_iqr_level = iqrdata[:2] | |
best_label = max_iqr[layer].max(0)[1] | |
best_score = lscore[best_label, torch.arange(lscore.shape[1])] | |
best_qscore = qscore[best_label, torch.arange(lscore.shape[1])] | |
else: | |
# Otherwise, assign labels based on max iou | |
best_score, best_label = lscore.max(0) | |
best_qscore = qscore[best_label, torch.arange(qscore.shape[1])] | |
record['iou_threshold'] = iou_threshold, | |
for u, urec in enumerate(units): | |
score, qscore, label = ( | |
best_score[u], best_qscore[u], best_label[u]) | |
urec.update(dict( | |
iou=score.item(), | |
iou_iqr=qscore.item(), | |
lc=lcs[label].item(), | |
cc=cc[catnumber[labelnames[label][1]], u].item(), | |
ic=ic[label, u].item(), | |
interp=(qscore.item() > iqr_threshold and | |
score.item() > iou_threshold), | |
iou_labelnum=label.item(), | |
iou_label=labelnames[label.item()][0], | |
iou_cat=labelnames[label.item()][1], | |
)) | |
if maxioudata is not None: | |
max_iou, max_iou_level, max_iou_quantile = maxioudata | |
qualified_iou = max_iou[layer].clone() | |
# qualified_iou[max_iou_quantile[layer] > 0.75] = 0 | |
best_score, best_label = qualified_iou.max(0) | |
for u, urec in enumerate(units): | |
urec.update(dict( | |
maxiou=best_score[u].item(), | |
maxiou_label=labelnames[best_label[u].item()][0], | |
maxiou_cat=labelnames[best_label[u].item()][1], | |
maxiou_level=max_iou_level[layer][best_label[u], u].item(), | |
maxiou_quantile=max_iou_quantile[layer][ | |
best_label[u], u].item())) | |
if iqrdata is not None: | |
[max_iqr, max_iqr_level, max_iqr_quantile, | |
max_iqr_iou, max_iqr_agreement] = iqrdata | |
qualified_iqr = max_iqr[layer].clone() | |
qualified_iqr[max_iqr_quantile[layer] > 0.5] = 0 | |
best_score, best_label = qualified_iqr.max(0) | |
for u, urec in enumerate(units): | |
urec.update(dict( | |
iqr=best_score[u].item(), | |
iqr_label=labelnames[best_label[u].item()][0], | |
iqr_cat=labelnames[best_label[u].item()][1], | |
iqr_level=max_iqr_level[layer][best_label[u], u].item(), | |
iqr_quantile=max_iqr_quantile[layer][ | |
best_label[u], u].item(), | |
iqr_iou=max_iqr_iou[layer][best_label[u], u].item() | |
)) | |
if covariancedata is not None: | |
score = covariancedata[layer].correlation() | |
best_score, best_label = score.max(1) | |
for u, urec in enumerate(units): | |
urec.update(dict( | |
cor=best_score[u].item(), | |
cor_label=labelnames[best_label[u].item()][0], | |
cor_cat=labelnames[best_label[u].item()][1] | |
)) | |
if mergedata is not None: | |
# Final step: if the user passed any data to merge into the | |
# units, merge them now. This can be used, for example, to | |
# indiate that a unit is not interpretable based on some | |
# outside analysis of unit statistics. | |
for lrec in mergedata.get('layers', []): | |
if lrec['layer'] == layer: | |
break | |
else: | |
lrec = None | |
for u, urec in enumerate(lrec.get('units', []) if lrec else []): | |
units[u].update(urec) | |
# After populating per-unit info, populate per-layer ranking info | |
if labeldata is not None: | |
# Collect all labeled units | |
labelunits = defaultdict(list) | |
all_labelunits = defaultdict(list) | |
for u, urec in enumerate(units): | |
if urec['interp']: | |
labelunits[urec['iou_labelnum']].append(u) | |
all_labelunits[urec['iou_labelnum']].append(u) | |
# Sort all units in order with most popular label first. | |
label_ordering = sorted(units, | |
# Sort by: | |
key=lambda r: (-1 if r['interp'] else 0, # interpretable | |
-len(labelunits[r['iou_labelnum']]), # label freq, score | |
-max([units[u]['iou'] | |
for u in labelunits[r['iou_labelnum']]], default=0), | |
r['iou_labelnum'], # label | |
-r['iou'])) # unit score | |
# Add label and iou ranking. | |
rankings.append(dict(name="label", score=(numpy.argsort(list( | |
ur['unit'] for ur in label_ordering))).tolist())) | |
rankings.append(dict(name="max iou", metric="iou", score=list( | |
-ur['iou'] for ur in units))) | |
# Add ranking for top labels | |
# for labelnum in [n for n in sorted( | |
# all_labelunits.keys(), key=lambda x: | |
# -len(all_labelunits[x])) if len(all_labelunits[n])]: | |
# label = labelnames[labelnum][0] | |
# rankings.append(dict(name="%s-iou" % label, | |
# concept=label, metric='iou', | |
# score=(-lscore[labelnum, :]).tolist())) | |
# Collate labels by category then frequency. | |
record['labels'] = [dict( | |
label=labelnames[label][0], | |
labelnum=label, | |
units=labelunits[label], | |
cat=labelnames[label][1]) | |
for label in (sorted(labelunits.keys(), | |
# Sort by: | |
key=lambda l: (catorder.get( # category | |
labelnames[l][1], 0), | |
-len(labelunits[l]), # label freq | |
-max([units[u]['iou'] for u in labelunits[l]], | |
default=0) # score | |
))) if len(labelunits[label])] | |
# Total number of interpretable units. | |
record['interpretable'] = sum(len(group['units']) | |
for group in record['labels']) | |
# Make a bargraph of labels | |
os.makedirs(os.path.join(outdir, safe_dir_name(layer)), | |
exist_ok=True) | |
catgroups = OrderedDict() | |
for _, cat in sorted([(v, k) for k, v in catorder.items()]): | |
catgroups[cat] = [] | |
for rec in record['labels']: | |
if rec['cat'] not in catgroups: | |
catgroups[rec['cat']] = [] | |
catgroups[rec['cat']].append(rec['label']) | |
make_svg_bargraph( | |
[rec['label'] for rec in record['labels']], | |
[len(rec['units']) for rec in record['labels']], | |
[(cat, len(group)) for cat, group in catgroups.items()], | |
filename=os.path.join(outdir, safe_dir_name(layer), | |
'bargraph.svg')) | |
# Only show the bargraph if it is non-empty. | |
if len(record['labels']): | |
record['bargraph'] = 'bargraph.svg' | |
if maxioudata is not None: | |
rankings.append(dict(name="max maxiou", metric="maxiou", score=list( | |
-ur['maxiou'] for ur in units))) | |
if iqrdata is not None: | |
rankings.append(dict(name="max iqr", metric="iqr", score=list( | |
-ur['iqr'] for ur in units))) | |
if covariancedata is not None: | |
rankings.append(dict(name="max cor", metric="cor", score=list( | |
-ur['cor'] for ur in units))) | |
all_layers.append(record) | |
# Now add the same rankings to every layer... | |
all_labels = None | |
if rank_all_labels: | |
all_labels = [name for name, cat in labelnames] | |
if labeldata is not None: | |
# Count layers+quadrants with a given label, and sort by freq | |
counted_labels = defaultdict(int) | |
for label in [ | |
re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '', unitrec['iou_label']) | |
for record in all_layers for unitrec in record['units']]: | |
counted_labels[label] += 1 | |
if all_labels is None: | |
all_labels = [label for count, label in sorted((-v, k) | |
for k, v in counted_labels.items())] | |
for record in all_layers: | |
layer = record['layer'] | |
for label in all_labels: | |
labelnum = labelnumber[label] | |
record['rankings'].append(dict(name="%s-iou" % label, | |
concept=label, metric='iou', | |
score=(-iou_scores[layer][labelnum, :]).tolist())) | |
if maxioudata is not None: | |
if all_labels is None: | |
counted_labels = defaultdict(int) | |
for label in [ | |
re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '', | |
unitrec['maxiou_label']) | |
for record in all_layers for unitrec in record['units']]: | |
counted_labels[label] += 1 | |
all_labels = [label for count, label in sorted((-v, k) | |
for k, v in counted_labels.items())] | |
qualified_iou = max_iou[layer].clone() | |
qualified_iou[max_iou_quantile[layer] > 0.5] = 0 | |
for record in all_layers: | |
layer = record['layer'] | |
for label in all_labels: | |
labelnum = labelnumber[label] | |
record['rankings'].append(dict(name="%s-maxiou" % label, | |
concept=label, metric='maxiou', | |
score=(-qualified_iou[labelnum, :]).tolist())) | |
if iqrdata is not None: | |
if all_labels is None: | |
counted_labels = defaultdict(int) | |
for label in [ | |
re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '', | |
unitrec['iqr_label']) | |
for record in all_layers for unitrec in record['units']]: | |
counted_labels[label] += 1 | |
all_labels = [label for count, label in sorted((-v, k) | |
for k, v in counted_labels.items())] | |
# qualified_iqr[max_iqr_quantile[layer] > 0.5] = 0 | |
for record in all_layers: | |
layer = record['layer'] | |
qualified_iqr = max_iqr[layer].clone() | |
for label in all_labels: | |
labelnum = labelnumber[label] | |
record['rankings'].append(dict(name="%s-iqr" % label, | |
concept=label, metric='iqr', | |
score=(-qualified_iqr[labelnum, :]).tolist())) | |
if covariancedata is not None: | |
if all_labels is None: | |
counted_labels = defaultdict(int) | |
for label in [ | |
re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '', | |
unitrec['cor_label']) | |
for record in all_layers for unitrec in record['units']]: | |
counted_labels[label] += 1 | |
all_labels = [label for count, label in sorted((-v, k) | |
for k, v in counted_labels.items())] | |
for record in all_layers: | |
layer = record['layer'] | |
score = covariancedata[layer].correlation() | |
for label in all_labels: | |
labelnum = labelnumber[label] | |
record['rankings'].append(dict(name="%s-cor" % label, | |
concept=label, metric='cor', | |
score=(-score[:, labelnum]).tolist())) | |
for record in all_layers: | |
layer = record['layer'] | |
# Dump per-layer json inside per-layer directory | |
record['dirname'] = '.' | |
with open(os.path.join(outdir, safe_dir_name(layer), 'dissect.json'), | |
'w') as jsonfile: | |
top_record['layers'] = [record] | |
json.dump(top_record, jsonfile, indent=1) | |
# Copy the per-layer html | |
shutil.copy(os.path.join(srcdir, 'dissect.html'), | |
os.path.join(outdir, safe_dir_name(layer), 'dissect.html')) | |
record['dirname'] = safe_dir_name(layer) | |
# Dump all-layer json in parent directory | |
with open(os.path.join(outdir, 'dissect.json'), 'w') as jsonfile: | |
top_record['layers'] = all_layers | |
json.dump(top_record, jsonfile, indent=1) | |
# Copy the all-layer html | |
shutil.copy(os.path.join(srcdir, 'dissect.html'), | |
os.path.join(outdir, 'dissect.html')) | |
shutil.copy(os.path.join(srcdir, 'edit.html'), | |
os.path.join(outdir, 'edit.html')) | |
def generate_images(outdir, model, dataset, topk, levels, | |
segrunner, row_length=None, gap_pixels=5, | |
row_images=True, single_images=False, prefix='', | |
batch_size=100, num_workers=24): | |
''' | |
Creates an image strip file for every unit of every retained layer | |
of the model, in the format [outdir]/[layername]/[unitnum]-top.jpg. | |
Assumes that the indexes of topk refer to the indexes of dataset. | |
Limits each strip to the top row_length images. | |
''' | |
progress = default_progress() | |
needed_images = {} | |
if row_images is False: | |
row_length = 1 | |
# Pass 1: needed_images lists all images that are topk for some unit. | |
for layer in topk: | |
topresult = topk[layer].result()[1].cpu() | |
for unit, row in enumerate(topresult): | |
for rank, imgnum in enumerate(row[:row_length]): | |
imgnum = imgnum.item() | |
if imgnum not in needed_images: | |
needed_images[imgnum] = [] | |
needed_images[imgnum].append((layer, unit, rank)) | |
levels = {k: v.cpu().numpy() for k, v in levels.items()} | |
row_length = len(row[:row_length]) | |
needed_sample = FixedSubsetSampler(sorted(needed_images.keys())) | |
device = next(model.parameters()).device | |
segloader = torch.utils.data.DataLoader(dataset, | |
batch_size=batch_size, num_workers=num_workers, | |
pin_memory=(device.type == 'cuda'), | |
sampler=needed_sample) | |
vizgrid, maskgrid, origrid, seggrid = [{} for _ in range(4)] | |
# Pass 2: populate vizgrid with visualizations of top units. | |
pool = None | |
for i, batch in enumerate( | |
progress(segloader, desc='Making images')): | |
# Reverse transformation to get the image in byte form. | |
seg, _, byte_im, _ = segrunner.run_and_segment_batch(batch, model, | |
want_rgb=True) | |
torch_features = model.retained_features() | |
scale_offset = getattr(model, 'scale_offset', None) | |
if pool is None: | |
# Distribute the work across processes: create shared mmaps. | |
for layer, tf in torch_features.items(): | |
[vizgrid[layer], maskgrid[layer], origrid[layer], | |
seggrid[layer]] = [ | |
create_temp_mmap_grid((tf.shape[1], | |
byte_im.shape[1], row_length, | |
byte_im.shape[2] + gap_pixels, depth), | |
dtype='uint8', | |
fill=255) | |
for depth in [3, 4, 3, 3]] | |
# Pass those mmaps to worker processes. | |
pool = WorkerPool(worker=VisualizeImageWorker, | |
memmap_grid_info=[ | |
{layer: (g.filename, g.shape, g.dtype) | |
for layer, g in grid.items()} | |
for grid in [vizgrid, maskgrid, origrid, seggrid]]) | |
byte_im = byte_im.cpu().numpy() | |
numpy_seg = seg.cpu().numpy() | |
features = {} | |
for index in range(len(byte_im)): | |
imgnum = needed_sample.samples[index + i*segloader.batch_size] | |
for layer, unit, rank in needed_images[imgnum]: | |
if layer not in features: | |
features[layer] = torch_features[layer].cpu().numpy() | |
pool.add(layer, unit, rank, | |
byte_im[index], | |
features[layer][index, unit], | |
levels[layer][unit], | |
scale_offset[layer] if scale_offset else None, | |
numpy_seg[index]) | |
pool.join() | |
# Pass 3: save image strips as [outdir]/[layer]/[unitnum]-[top/orig].jpg | |
pool = WorkerPool(worker=SaveImageWorker) | |
for layer, vg in progress(vizgrid.items(), desc='Saving images'): | |
os.makedirs(os.path.join(outdir, safe_dir_name(layer), | |
prefix + 'image'), exist_ok=True) | |
if single_images: | |
os.makedirs(os.path.join(outdir, safe_dir_name(layer), | |
prefix + 's-image'), exist_ok=True) | |
og, sg, mg = origrid[layer], seggrid[layer], maskgrid[layer] | |
for unit in progress(range(len(vg)), desc='Units'): | |
for suffix, grid in [('top.jpg', vg), ('orig.jpg', og), | |
('seg.png', sg), ('mask.png', mg)]: | |
strip = grid[unit].reshape( | |
(grid.shape[1], grid.shape[2] * grid.shape[3], | |
grid.shape[4])) | |
if row_images: | |
filename = os.path.join(outdir, safe_dir_name(layer), | |
prefix + 'image', '%d-%s' % (unit, suffix)) | |
pool.add(strip[:,:-gap_pixels,:].copy(), filename) | |
# Image.fromarray(strip[:,:-gap_pixels,:]).save(filename, | |
# optimize=True, quality=80) | |
if single_images: | |
single_filename = os.path.join(outdir, safe_dir_name(layer), | |
prefix + 's-image', '%d-%s' % (unit, suffix)) | |
pool.add(strip[:,:strip.shape[1] // row_length | |
- gap_pixels,:].copy(), single_filename) | |
# Image.fromarray(strip[:,:strip.shape[1] // row_length | |
# - gap_pixels,:]).save(single_filename, | |
# optimize=True, quality=80) | |
pool.join() | |
# Delete the shared memory map files | |
clear_global_shared_files([g.filename | |
for grid in [vizgrid, maskgrid, origrid, seggrid] | |
for g in grid.values()]) | |
global_shared_files = {} | |
def create_temp_mmap_grid(shape, dtype, fill): | |
dtype = numpy.dtype(dtype) | |
filename = os.path.join(tempfile.mkdtemp(), 'temp-%s-%s.mmap' % | |
('x'.join('%d' % s for s in shape), dtype.name)) | |
fid = open(filename, mode='w+b') | |
original = numpy.memmap(fid, dtype=dtype, mode='w+', shape=shape) | |
original.fid = fid | |
original[...] = fill | |
global_shared_files[filename] = original | |
return original | |
def shared_temp_mmap_grid(filename, shape, dtype): | |
if filename not in global_shared_files: | |
global_shared_files[filename] = numpy.memmap( | |
filename, dtype=dtype, mode='r+', shape=shape) | |
return global_shared_files[filename] | |
def clear_global_shared_files(filenames): | |
for fn in filenames: | |
if fn in global_shared_files: | |
del global_shared_files[fn] | |
try: | |
os.unlink(fn) | |
except OSError: | |
pass | |
class VisualizeImageWorker(WorkerBase): | |
def setup(self, memmap_grid_info): | |
self.vizgrid, self.maskgrid, self.origrid, self.seggrid = [ | |
{layer: shared_temp_mmap_grid(*info) | |
for layer, info in grid.items()} | |
for grid in memmap_grid_info] | |
def work(self, layer, unit, rank, | |
byte_im, acts, level, scale_offset, seg): | |
self.origrid[layer][unit,:,rank,:byte_im.shape[0],:] = byte_im | |
[self.vizgrid[layer][unit,:,rank,:byte_im.shape[0],:], | |
self.maskgrid[layer][unit,:,rank,:byte_im.shape[0],:]] = ( | |
activation_visualization( | |
byte_im, | |
acts, | |
level, | |
scale_offset=scale_offset, | |
return_mask=True)) | |
self.seggrid[layer][unit,:,rank,:byte_im.shape[0],:] = ( | |
segment_visualization(seg, byte_im.shape[0:2])) | |
class SaveImageWorker(WorkerBase): | |
def work(self, data, filename): | |
Image.fromarray(data).save(filename, optimize=True, quality=80) | |
def score_tally_stats(label_category, tc, truth, cc, ic): | |
pred = cc[label_category] | |
total = tc[label_category][:, None] | |
truth = truth[:, None] | |
epsilon = 1e-20 # avoid division-by-zero | |
union = pred + truth - ic | |
iou = ic.double() / (union.double() + epsilon) | |
arr = torch.empty(size=(2, 2) + ic.shape, dtype=ic.dtype, device=ic.device) | |
arr[0, 0] = ic | |
arr[0, 1] = pred - ic | |
arr[1, 0] = truth - ic | |
arr[1, 1] = total - union | |
arr = arr.double() / total.double() | |
mi = mutual_information(arr) | |
je = joint_entropy(arr) | |
iqr = mi / je | |
iqr[torch.isnan(iqr)] = 0 # Zero out any 0/0 | |
return iou, iqr | |
def collect_quantiles_and_topk(outdir, model, segloader, | |
segrunner, k=100, resolution=1024): | |
''' | |
Collects (estimated) quantile information and (exact) sorted top-K lists | |
for every channel in the retained layers of the model. Returns | |
a map of quantiles (one RunningQuantile for each layer) along with | |
a map of topk (one RunningTopK for each layer). | |
''' | |
device = next(model.parameters()).device | |
features = model.retained_features() | |
cached_quantiles = { | |
layer: load_quantile_if_present(os.path.join(outdir, | |
safe_dir_name(layer)), 'quantiles.npz', | |
device=torch.device('cpu')) | |
for layer in features } | |
cached_topks = { | |
layer: load_topk_if_present(os.path.join(outdir, | |
safe_dir_name(layer)), 'topk.npz', | |
device=torch.device('cpu')) | |
for layer in features } | |
if (all(value is not None for value in cached_quantiles.values()) and | |
all(value is not None for value in cached_topks.values())): | |
return cached_quantiles, cached_topks | |
layer_batch_size = 8 | |
all_layers = list(features.keys()) | |
layer_batches = [all_layers[i:i+layer_batch_size] | |
for i in range(0, len(all_layers), layer_batch_size)] | |
quantiles, topks = {}, {} | |
progress = default_progress() | |
for layer_batch in layer_batches: | |
for i, batch in enumerate(progress(segloader, desc='Quantiles')): | |
# We don't actually care about the model output. | |
model(batch[0].to(device)) | |
features = model.retained_features() | |
# We care about the retained values | |
for key in layer_batch: | |
value = features[key] | |
if topks.get(key, None) is None: | |
topks[key] = RunningTopK(k) | |
if quantiles.get(key, None) is None: | |
quantiles[key] = RunningQuantile(resolution=resolution) | |
topvalue = value | |
if len(value.shape) > 2: | |
topvalue, _ = value.view(*(value.shape[:2] + (-1,))).max(2) | |
# Put the channel index last. | |
value = value.permute( | |
(0,) + tuple(range(2, len(value.shape))) + (1,) | |
).contiguous().view(-1, value.shape[1]) | |
quantiles[key].add(value) | |
topks[key].add(topvalue) | |
# Save GPU memory | |
for key in layer_batch: | |
quantiles[key].to_(torch.device('cpu')) | |
topks[key].to_(torch.device('cpu')) | |
for layer in quantiles: | |
save_state_dict(quantiles[layer], | |
os.path.join(outdir, safe_dir_name(layer), 'quantiles.npz')) | |
save_state_dict(topks[layer], | |
os.path.join(outdir, safe_dir_name(layer), 'topk.npz')) | |
return quantiles, topks | |
def collect_bincounts(outdir, model, segloader, levels, segrunner): | |
''' | |
Returns label_counts, category_activation_counts, and intersection_counts, | |
across the data set, counting the pixels of intersection between upsampled, | |
thresholded model featuremaps, with segmentation classes in the segloader. | |
label_counts (independent of model): pixels across the data set that | |
are labeled with the given label. | |
category_activation_counts (one per layer): for each feature channel, | |
pixels across the dataset where the channel exceeds the level | |
threshold. There is one count per category: activations only | |
contribute to the categories for which any category labels are | |
present on the images. | |
intersection_counts (one per layer): for each feature channel and | |
label, pixels across the dataset where the channel exceeds | |
the level, and the labeled segmentation class is also present. | |
This is a performance-sensitive function. Best performance is | |
achieved with a counting scheme which assumes a segloader with | |
batch_size 1. | |
''' | |
# Load cached data if present | |
(iou_scores, iqr_scores, | |
total_counts, label_counts, category_activation_counts, | |
intersection_counts) = {}, {}, None, None, {}, {} | |
found_all = True | |
for layer in model.retained_features(): | |
filename = os.path.join(outdir, safe_dir_name(layer), 'bincounts.npz') | |
if os.path.isfile(filename): | |
data = numpy.load(filename) | |
iou_scores[layer] = torch.from_numpy(data['iou_scores']) | |
iqr_scores[layer] = torch.from_numpy(data['iqr_scores']) | |
total_counts = torch.from_numpy(data['total_counts']) | |
label_counts = torch.from_numpy(data['label_counts']) | |
category_activation_counts[layer] = torch.from_numpy( | |
data['category_activation_counts']) | |
intersection_counts[layer] = torch.from_numpy( | |
data['intersection_counts']) | |
else: | |
found_all = False | |
if found_all: | |
return (iou_scores, iqr_scores, | |
total_counts, label_counts, category_activation_counts, | |
intersection_counts) | |
device = next(model.parameters()).device | |
labelcat, categories = segrunner.get_label_and_category_names() | |
label_category = [categories.index(c) if c in categories else 0 | |
for l, c in labelcat] | |
num_labels, num_categories = (len(n) for n in [labelcat, categories]) | |
# One-hot vector of category for each label | |
labelcat = torch.zeros(num_labels, num_categories, | |
dtype=torch.long, device=device) | |
labelcat.scatter_(1, torch.from_numpy(numpy.array(label_category, | |
dtype='int64')).to(device)[:,None], 1) | |
# Running bincounts | |
# activation_counts = {} | |
assert segloader.batch_size == 1 # category_activation_counts needs this. | |
category_activation_counts = {} | |
intersection_counts = {} | |
label_counts = torch.zeros(num_labels, dtype=torch.long, device=device) | |
total_counts = torch.zeros(num_categories, dtype=torch.long, device=device) | |
progress = default_progress() | |
scale_offset_map = getattr(model, 'scale_offset', None) | |
upsample_grids = {} | |
# total_batch_categories = torch.zeros( | |
# labelcat.shape[1], dtype=torch.long, device=device) | |
for i, batch in enumerate(progress(segloader, desc='Bincounts')): | |
seg, batch_label_counts, _, imshape = segrunner.run_and_segment_batch( | |
batch, model, want_bincount=True, want_rgb=True) | |
bc = batch_label_counts.cpu() | |
batch_label_counts = batch_label_counts.to(device) | |
seg = seg.to(device) | |
features = model.retained_features() | |
# Accumulate bincounts and identify nonzeros | |
label_counts += batch_label_counts[0] | |
batch_labels = bc[0].nonzero()[:,0] | |
batch_categories = labelcat[batch_labels].max(0)[0] | |
total_counts += batch_categories * ( | |
seg.shape[0] * seg.shape[2] * seg.shape[3]) | |
for key, value in features.items(): | |
if key not in upsample_grids: | |
upsample_grids[key] = upsample_grid(value.shape[2:], | |
seg.shape[2:], imshape, | |
scale_offset=scale_offset_map.get(key, None) | |
if scale_offset_map is not None else None, | |
dtype=value.dtype, device=value.device) | |
upsampled = torch.nn.functional.grid_sample(value, | |
upsample_grids[key], padding_mode='border') | |
amask = (upsampled > levels[key][None,:,None,None].to( | |
upsampled.device)) | |
ac = amask.int().view(amask.shape[1], -1).sum(1) | |
# if key not in activation_counts: | |
# activation_counts[key] = ac | |
# else: | |
# activation_counts[key] += ac | |
# The fastest approach: sum over each label separately! | |
for label in batch_labels.tolist(): | |
if label == 0: | |
continue # ignore the background label | |
imask = amask * ((seg == label).max(dim=1, keepdim=True)[0]) | |
ic = imask.int().view(imask.shape[1], -1).sum(1) | |
if key not in intersection_counts: | |
intersection_counts[key] = torch.zeros(num_labels, | |
amask.shape[1], dtype=torch.long, device=device) | |
intersection_counts[key][label] += ic | |
# Count activations within images that have category labels. | |
# Note: This only makes sense with batch-size one | |
# total_batch_categories += batch_categories | |
cc = batch_categories[:,None] * ac[None,:] | |
if key not in category_activation_counts: | |
category_activation_counts[key] = cc | |
else: | |
category_activation_counts[key] += cc | |
iou_scores = {} | |
iqr_scores = {} | |
for k in intersection_counts: | |
iou_scores[k], iqr_scores[k] = score_tally_stats( | |
label_category, total_counts, label_counts, | |
category_activation_counts[k], intersection_counts[k]) | |
for k in intersection_counts: | |
numpy.savez(os.path.join(outdir, safe_dir_name(k), 'bincounts.npz'), | |
iou_scores=iou_scores[k].cpu().numpy(), | |
iqr_scores=iqr_scores[k].cpu().numpy(), | |
total_counts=total_counts.cpu().numpy(), | |
label_counts=label_counts.cpu().numpy(), | |
category_activation_counts=category_activation_counts[k] | |
.cpu().numpy(), | |
intersection_counts=intersection_counts[k].cpu().numpy(), | |
levels=levels[k].cpu().numpy()) | |
return (iou_scores, iqr_scores, | |
total_counts, label_counts, category_activation_counts, | |
intersection_counts) | |
def collect_cond_quantiles(outdir, model, segloader, segrunner): | |
''' | |
Returns maxiou and maxiou_level across the data set, one per layer. | |
This is a performance-sensitive function. Best performance is | |
achieved with a counting scheme which assumes a segloader with | |
batch_size 1. | |
''' | |
device = next(model.parameters()).device | |
cached_cond_quantiles = { | |
layer: load_conditional_quantile_if_present(os.path.join(outdir, | |
safe_dir_name(layer)), 'cond_quantiles.npz') # on cpu | |
for layer in model.retained_features() } | |
label_fracs = load_npy_if_present(outdir, 'label_fracs.npy', 'cpu') | |
if label_fracs is not None and all( | |
value is not None for value in cached_cond_quantiles.values()): | |
return cached_cond_quantiles, label_fracs | |
labelcat, categories = segrunner.get_label_and_category_names() | |
label_category = [categories.index(c) if c in categories else 0 | |
for l, c in labelcat] | |
num_labels, num_categories = (len(n) for n in [labelcat, categories]) | |
# One-hot vector of category for each label | |
labelcat = torch.zeros(num_labels, num_categories, | |
dtype=torch.long, device=device) | |
labelcat.scatter_(1, torch.from_numpy(numpy.array(label_category, | |
dtype='int64')).to(device)[:,None], 1) | |
# Running maxiou | |
assert segloader.batch_size == 1 # category_activation_counts needs this. | |
conditional_quantiles = {} | |
label_counts = torch.zeros(num_labels, dtype=torch.long, device=device) | |
pixel_count = 0 | |
progress = default_progress() | |
scale_offset_map = getattr(model, 'scale_offset', None) | |
upsample_grids = {} | |
common_conditions = set() | |
if label_fracs is None or label_fracs is 0: | |
for i, batch in enumerate(progress(segloader, desc='label fracs')): | |
seg, batch_label_counts, im, _ = segrunner.run_and_segment_batch( | |
batch, model, want_bincount=True, want_rgb=True) | |
batch_label_counts = batch_label_counts.to(device) | |
features = model.retained_features() | |
# Accumulate bincounts and identify nonzeros | |
label_counts += batch_label_counts[0] | |
pixel_count += seg.shape[2] * seg.shape[3] | |
label_fracs = (label_counts.cpu().float() / pixel_count)[:, None, None] | |
numpy.save(os.path.join(outdir, 'label_fracs.npy'), label_fracs) | |
skip_threshold = 1e-4 | |
skip_labels = set(i.item() | |
for i in (label_fracs.view(-1) < skip_threshold).nonzero().view(-1)) | |
for layer in progress(model.retained_features().keys(), desc='CQ layers'): | |
if cached_cond_quantiles.get(layer, None) is not None: | |
conditional_quantiles[layer] = cached_cond_quantiles[layer] | |
continue | |
for i, batch in enumerate(progress(segloader, desc='Condquant')): | |
seg, batch_label_counts, _, imshape = ( | |
segrunner.run_and_segment_batch( | |
batch, model, want_bincount=True, want_rgb=True)) | |
bc = batch_label_counts.cpu() | |
batch_label_counts = batch_label_counts.to(device) | |
features = model.retained_features() | |
# Accumulate bincounts and identify nonzeros | |
label_counts += batch_label_counts[0] | |
pixel_count += seg.shape[2] * seg.shape[3] | |
batch_labels = bc[0].nonzero()[:,0] | |
batch_categories = labelcat[batch_labels].max(0)[0] | |
cpu_seg = None | |
value = features[layer] | |
if layer not in upsample_grids: | |
upsample_grids[layer] = upsample_grid(value.shape[2:], | |
seg.shape[2:], imshape, | |
scale_offset=scale_offset_map.get(layer, None) | |
if scale_offset_map is not None else None, | |
dtype=value.dtype, device=value.device) | |
if layer not in conditional_quantiles: | |
conditional_quantiles[layer] = RunningConditionalQuantile( | |
resolution=2048) | |
upsampled = torch.nn.functional.grid_sample(value, | |
upsample_grids[layer], padding_mode='border').view( | |
value.shape[1], -1) | |
conditional_quantiles[layer].add(('all',), upsampled.t()) | |
cpu_upsampled = None | |
for label in batch_labels.tolist(): | |
if label in skip_labels: | |
continue | |
label_key = ('label', label) | |
if label_key in common_conditions: | |
imask = (seg == label).max(dim=1)[0].view(-1) | |
intersected = upsampled[:, imask] | |
conditional_quantiles[layer].add(('label', label), | |
intersected.t()) | |
else: | |
if cpu_seg is None: | |
cpu_seg = seg.cpu() | |
if cpu_upsampled is None: | |
cpu_upsampled = upsampled.cpu() | |
imask = (cpu_seg == label).max(dim=1)[0].view(-1) | |
intersected = cpu_upsampled[:, imask] | |
conditional_quantiles[layer].add(('label', label), | |
intersected.t()) | |
if num_categories > 1: | |
for cat in batch_categories.nonzero()[:,0]: | |
conditional_quantiles[layer].add(('cat', cat.item()), | |
upsampled.t()) | |
# Move the most common conditions to the GPU. | |
if i and not i & (i - 1): # if i is a power of 2: | |
cq = conditional_quantiles[layer] | |
common_conditions = set(cq.most_common_conditions(64)) | |
cq.to_('cpu', [k for k in cq.running_quantiles.keys() | |
if k not in common_conditions]) | |
# When a layer is done, get it off the GPU | |
conditional_quantiles[layer].to_('cpu') | |
label_fracs = (label_counts.cpu().float() / pixel_count)[:, None, None] | |
for cq in conditional_quantiles.values(): | |
cq.to_('cpu') | |
for layer in conditional_quantiles: | |
save_state_dict(conditional_quantiles[layer], | |
os.path.join(outdir, safe_dir_name(layer), 'cond_quantiles.npz')) | |
numpy.save(os.path.join(outdir, 'label_fracs.npy'), label_fracs) | |
return conditional_quantiles, label_fracs | |
def collect_maxiou(outdir, model, segloader, segrunner): | |
''' | |
Returns maxiou and maxiou_level across the data set, one per layer. | |
This is a performance-sensitive function. Best performance is | |
achieved with a counting scheme which assumes a segloader with | |
batch_size 1. | |
''' | |
device = next(model.parameters()).device | |
conditional_quantiles, label_fracs = collect_cond_quantiles( | |
outdir, model, segloader, segrunner) | |
labelcat, categories = segrunner.get_label_and_category_names() | |
label_category = [categories.index(c) if c in categories else 0 | |
for l, c in labelcat] | |
num_labels, num_categories = (len(n) for n in [labelcat, categories]) | |
label_list = [('label', i) for i in range(num_labels)] | |
category_list = [('all',)] if num_categories <= 1 else ( | |
[('cat', i) for i in range(num_categories)]) | |
max_iou, max_iou_level, max_iou_quantile = {}, {}, {} | |
fracs = torch.logspace(-3, 0, 100) | |
progress = default_progress() | |
for layer, cq in progress(conditional_quantiles.items(), desc='Maxiou'): | |
levels = cq.conditional(('all',)).quantiles(1 - fracs) | |
denoms = 1 - cq.collected_normalize(category_list, levels) | |
isects = (1 - cq.collected_normalize(label_list, levels)) * label_fracs | |
unions = label_fracs + denoms[label_category, :, :] - isects | |
iou = isects / unions | |
# TODO: erase any for which threshold is bad | |
max_iou[layer], level_bucket = iou.max(2) | |
max_iou_level[layer] = levels[ | |
torch.arange(levels.shape[0])[None,:], level_bucket] | |
max_iou_quantile[layer] = fracs[level_bucket] | |
for layer in model.retained_features(): | |
numpy.savez(os.path.join(outdir, safe_dir_name(layer), 'max_iou.npz'), | |
max_iou=max_iou[layer].cpu().numpy(), | |
max_iou_level=max_iou_level[layer].cpu().numpy(), | |
max_iou_quantile=max_iou_quantile[layer].cpu().numpy()) | |
return (max_iou, max_iou_level, max_iou_quantile) | |
def collect_iqr(outdir, model, segloader, segrunner): | |
''' | |
Returns iqr and iqr_level. | |
This is a performance-sensitive function. Best performance is | |
achieved with a counting scheme which assumes a segloader with | |
batch_size 1. | |
''' | |
max_iqr, max_iqr_level, max_iqr_quantile, max_iqr_iou = {}, {}, {}, {} | |
max_iqr_agreement = {} | |
found_all = True | |
for layer in model.retained_features(): | |
filename = os.path.join(outdir, safe_dir_name(layer), 'iqr.npz') | |
if os.path.isfile(filename): | |
data = numpy.load(filename) | |
max_iqr[layer] = torch.from_numpy(data['max_iqr']) | |
max_iqr_level[layer] = torch.from_numpy(data['max_iqr_level']) | |
max_iqr_quantile[layer] = torch.from_numpy(data['max_iqr_quantile']) | |
max_iqr_iou[layer] = torch.from_numpy(data['max_iqr_iou']) | |
max_iqr_agreement[layer] = torch.from_numpy( | |
data['max_iqr_agreement']) | |
else: | |
found_all = False | |
if found_all: | |
return (max_iqr, max_iqr_level, max_iqr_quantile, max_iqr_iou, | |
max_iqr_agreement) | |
device = next(model.parameters()).device | |
conditional_quantiles, label_fracs = collect_cond_quantiles( | |
outdir, model, segloader, segrunner) | |
labelcat, categories = segrunner.get_label_and_category_names() | |
label_category = [categories.index(c) if c in categories else 0 | |
for l, c in labelcat] | |
num_labels, num_categories = (len(n) for n in [labelcat, categories]) | |
label_list = [('label', i) for i in range(num_labels)] | |
category_list = [('all',)] if num_categories <= 1 else ( | |
[('cat', i) for i in range(num_categories)]) | |
full_mi, full_je, full_iqr = {}, {}, {} | |
fracs = torch.logspace(-3, 0, 100) | |
progress = default_progress() | |
for layer, cq in progress(conditional_quantiles.items(), desc='IQR'): | |
levels = cq.conditional(('all',)).quantiles(1 - fracs) | |
truth = label_fracs.to(device) | |
preds = (1 - cq.collected_normalize(category_list, levels) | |
)[label_category, :, :].to(device) | |
cond_isects = 1 - cq.collected_normalize(label_list, levels).to(device) | |
isects = cond_isects * truth | |
unions = truth + preds - isects | |
arr = torch.empty(size=(2, 2) + isects.shape, dtype=isects.dtype, | |
device=device) | |
arr[0, 0] = isects | |
arr[0, 1] = preds - isects | |
arr[1, 0] = truth - isects | |
arr[1, 1] = 1 - unions | |
arr.clamp_(0, 1) | |
mi = mutual_information(arr) | |
mi[:,:,-1] = 0 # at the 1.0 quantile should be no MI. | |
# Don't trust mi when less than label_frac is less than 1e-3, | |
# because our samples are too small. | |
mi[label_fracs.view(-1) < 1e-3, :, :] = 0 | |
je = joint_entropy(arr) | |
iqr = mi / je | |
iqr[torch.isnan(iqr)] = 0 # Zero out any 0/0 | |
full_mi[layer] = mi.cpu() | |
full_je[layer] = je.cpu() | |
full_iqr[layer] = iqr.cpu() | |
del mi, je | |
agreement = isects + arr[1, 1] | |
# When optimizing, maximize only over those pairs where the | |
# unit is positively correlated with the label, and where the | |
# threshold level is positive | |
positive_iqr = iqr | |
positive_iqr[agreement <= 0.8] = 0 | |
positive_iqr[(levels <= 0.0)[None, :, :].expand(positive_iqr.shape)] = 0 | |
# TODO: erase any for which threshold is bad | |
maxiqr, level_bucket = positive_iqr.max(2) | |
max_iqr[layer] = maxiqr.cpu() | |
max_iqr_level[layer] = levels.to(device)[ | |
torch.arange(levels.shape[0])[None,:], level_bucket].cpu() | |
max_iqr_quantile[layer] = fracs.to(device)[level_bucket].cpu() | |
max_iqr_agreement[layer] = agreement[ | |
torch.arange(agreement.shape[0])[:, None], | |
torch.arange(agreement.shape[1])[None, :], | |
level_bucket].cpu() | |
# Compute the iou that goes with each maximized iqr | |
matching_iou = (isects[ | |
torch.arange(isects.shape[0])[:, None], | |
torch.arange(isects.shape[1])[None, :], | |
level_bucket] / | |
unions[ | |
torch.arange(unions.shape[0])[:, None], | |
torch.arange(unions.shape[1])[None, :], | |
level_bucket]) | |
matching_iou[torch.isnan(matching_iou)] = 0 | |
max_iqr_iou[layer] = matching_iou.cpu() | |
for layer in model.retained_features(): | |
numpy.savez(os.path.join(outdir, safe_dir_name(layer), 'iqr.npz'), | |
max_iqr=max_iqr[layer].cpu().numpy(), | |
max_iqr_level=max_iqr_level[layer].cpu().numpy(), | |
max_iqr_quantile=max_iqr_quantile[layer].cpu().numpy(), | |
max_iqr_iou=max_iqr_iou[layer].cpu().numpy(), | |
max_iqr_agreement=max_iqr_agreement[layer].cpu().numpy(), | |
full_mi=full_mi[layer].cpu().numpy(), | |
full_je=full_je[layer].cpu().numpy(), | |
full_iqr=full_iqr[layer].cpu().numpy()) | |
return (max_iqr, max_iqr_level, max_iqr_quantile, max_iqr_iou, | |
max_iqr_agreement) | |
def mutual_information(arr): | |
total = 0 | |
for j in range(arr.shape[0]): | |
for k in range(arr.shape[1]): | |
joint = arr[j,k] | |
ind = arr[j,:].sum(dim=0) * arr[:,k].sum(dim=0) | |
term = joint * (joint / ind).log() | |
term[torch.isnan(term)] = 0 | |
total += term | |
return total.clamp_(0) | |
def joint_entropy(arr): | |
total = 0 | |
for j in range(arr.shape[0]): | |
for k in range(arr.shape[1]): | |
joint = arr[j,k] | |
term = joint * joint.log() | |
term[torch.isnan(term)] = 0 | |
total += term | |
return (-total).clamp_(0) | |
def information_quality_ratio(arr): | |
iqr = mutual_information(arr) / joint_entropy(arr) | |
iqr[torch.isnan(iqr)] = 0 | |
return iqr | |
def collect_covariance(outdir, model, segloader, segrunner): | |
''' | |
Returns label_mean, label_variance, unit_mean, unit_variance, | |
and cross_covariance across the data set. | |
label_mean, label_variance (independent of model): | |
treating the label as a one-hot, each label's mean and variance. | |
unit_mean, unit_variance (one per layer): for each feature channel, | |
the mean and variance of the activations in that channel. | |
cross_covariance (one per layer): the cross covariance between the | |
labels and the units in the layer. | |
''' | |
device = next(model.parameters()).device | |
cached_covariance = { | |
layer: load_covariance_if_present(os.path.join(outdir, | |
safe_dir_name(layer)), 'covariance.npz', device=device) | |
for layer in model.retained_features() } | |
if all(value is not None for value in cached_covariance.values()): | |
return cached_covariance | |
labelcat, categories = segrunner.get_label_and_category_names() | |
label_category = [categories.index(c) if c in categories else 0 | |
for l, c in labelcat] | |
num_labels, num_categories = (len(n) for n in [labelcat, categories]) | |
# Running covariance | |
cov = {} | |
progress = default_progress() | |
scale_offset_map = getattr(model, 'scale_offset', None) | |
upsample_grids = {} | |
for i, batch in enumerate(progress(segloader, desc='Covariance')): | |
seg, _, _, imshape = segrunner.run_and_segment_batch(batch, model, | |
want_rgb=True) | |
features = model.retained_features() | |
ohfeats = multilabel_onehot(seg, num_labels, ignore_index=0) | |
# Accumulate bincounts and identify nonzeros | |
for key, value in features.items(): | |
if key not in upsample_grids: | |
upsample_grids[key] = upsample_grid(value.shape[2:], | |
seg.shape[2:], imshape, | |
scale_offset=scale_offset_map.get(key, None) | |
if scale_offset_map is not None else None, | |
dtype=value.dtype, device=value.device) | |
upsampled = torch.nn.functional.grid_sample(value, | |
upsample_grids[key].expand( | |
(value.shape[0],) + upsample_grids[key].shape[1:]), | |
padding_mode='border') | |
if key not in cov: | |
cov[key] = RunningCrossCovariance() | |
cov[key].add(upsampled, ohfeats) | |
for layer in cov: | |
save_state_dict(cov[layer], | |
os.path.join(outdir, safe_dir_name(layer), 'covariance.npz')) | |
return cov | |
def multilabel_onehot(labels, num_labels, dtype=None, ignore_index=None): | |
''' | |
Converts a multilabel tensor into a onehot tensor. | |
The input labels is a tensor of shape (samples, multilabels, y, x). | |
The output is a tensor of shape (samples, num_labels, y, x). | |
If ignore_index is specified, labels with that index are ignored. | |
Each x in labels should be 0 <= x < num_labels, or x == ignore_index. | |
''' | |
assert ignore_index is None or ignore_index <= 0 | |
if dtype is None: | |
dtype = torch.float | |
device = labels.device | |
chans = num_labels + (-ignore_index if ignore_index else 0) | |
outshape = (labels.shape[0], chans) + labels.shape[2:] | |
result = torch.zeros(outshape, device=device, dtype=dtype) | |
if ignore_index and ignore_index < 0: | |
labels = labels + (-ignore_index) | |
result.scatter_(1, labels, 1) | |
if ignore_index and ignore_index < 0: | |
result = result[:, -ignore_index:] | |
elif ignore_index is not None: | |
result[:, ignore_index] = 0 | |
return result | |
def load_npy_if_present(outdir, filename, device): | |
filepath = os.path.join(outdir, filename) | |
if os.path.isfile(filepath): | |
data = numpy.load(filepath) | |
return torch.from_numpy(data).to(device) | |
return 0 | |
def load_npz_if_present(outdir, filename, varnames, device): | |
filepath = os.path.join(outdir, filename) | |
if os.path.isfile(filepath): | |
data = numpy.load(filepath) | |
numpy_result = [data[n] for n in varnames] | |
return tuple(torch.from_numpy(data).to(device) for data in numpy_result) | |
return None | |
def load_quantile_if_present(outdir, filename, device): | |
filepath = os.path.join(outdir, filename) | |
if os.path.isfile(filepath): | |
data = numpy.load(filepath) | |
result = RunningQuantile(state=data) | |
result.to_(device) | |
return result | |
return None | |
def load_conditional_quantile_if_present(outdir, filename): | |
filepath = os.path.join(outdir, filename) | |
if os.path.isfile(filepath): | |
data = numpy.load(filepath) | |
result = RunningConditionalQuantile(state=data) | |
return result | |
return None | |
def load_topk_if_present(outdir, filename, device): | |
filepath = os.path.join(outdir, filename) | |
if os.path.isfile(filepath): | |
data = numpy.load(filepath) | |
result = RunningTopK(state=data) | |
result.to_(device) | |
return result | |
return None | |
def load_covariance_if_present(outdir, filename, device): | |
filepath = os.path.join(outdir, filename) | |
if os.path.isfile(filepath): | |
data = numpy.load(filepath) | |
result = RunningCrossCovariance(state=data) | |
result.to_(device) | |
return result | |
return None | |
def save_state_dict(obj, filepath): | |
dirname = os.path.dirname(filepath) | |
os.makedirs(dirname, exist_ok=True) | |
dic = obj.state_dict() | |
numpy.savez(filepath, **dic) | |
def upsample_grid(data_shape, target_shape, input_shape=None, | |
scale_offset=None, dtype=torch.float, device=None): | |
'''Prepares a grid to use with grid_sample to upsample a batch of | |
features in data_shape to the target_shape. Can use scale_offset | |
and input_shape to center the grid in a nondefault way: scale_offset | |
maps feature pixels to input_shape pixels, and it is assumed that | |
the target_shape is a uniform downsampling of input_shape.''' | |
# Default is that nothing is resized. | |
if target_shape is None: | |
target_shape = data_shape | |
# Make a default scale_offset to fill the image if there isn't one | |
if scale_offset is None: | |
scale = tuple(float(ts) / ds | |
for ts, ds in zip(target_shape, data_shape)) | |
offset = tuple(0.5 * s - 0.5 for s in scale) | |
else: | |
scale, offset = (v for v in zip(*scale_offset)) | |
# Handle downsampling for different input vs target shape. | |
if input_shape is not None: | |
scale = tuple(s * (ts - 1) / (ns - 1) | |
for s, ns, ts in zip(scale, input_shape, target_shape)) | |
offset = tuple(o * (ts - 1) / (ns - 1) | |
for o, ns, ts in zip(offset, input_shape, target_shape)) | |
# Pytorch needs target coordinates in terms of source coordinates [-1..1] | |
ty, tx = (((torch.arange(ts, dtype=dtype, device=device) - o) | |
* (2 / (s * (ss - 1))) - 1) | |
for ts, ss, s, o, in zip(target_shape, data_shape, scale, offset)) | |
# Whoa, note that grid_sample reverses the order y, x -> x, y. | |
grid = torch.stack( | |
(tx[None,:].expand(target_shape), ty[:,None].expand(target_shape)),2 | |
)[None,:,:,:].expand((1, target_shape[0], target_shape[1], 2)) | |
return grid | |
def safe_dir_name(filename): | |
keepcharacters = (' ','.','_','-') | |
return ''.join(c | |
for c in filename if c.isalnum() or c in keepcharacters).rstrip() | |
bargraph_palette = [ | |
('#4B4CBF', '#B6B6F2'), | |
('#55B05B', '#B6F2BA'), | |
('#50BDAC', '#A5E5DB'), | |
('#81C679', '#C0FF9B'), | |
('#F0883B', '#F2CFB6'), | |
('#D4CF24', '#F2F1B6'), | |
('#D92E2B', '#F2B6B6'), | |
('#AB6BC6', '#CFAAFF'), | |
] | |
def make_svg_bargraph(labels, heights, categories, | |
barheight=100, barwidth=12, show_labels=True, filename=None): | |
# if len(labels) == 0: | |
# return # Nothing to do | |
unitheight = float(barheight) / max(max(heights, default=1), 1) | |
textheight = barheight if show_labels else 0 | |
labelsize = float(barwidth) | |
gap = float(barwidth) / 4 | |
textsize = barwidth + gap | |
rollup = max(heights, default=1) | |
textmargin = float(labelsize) * 2 / 3 | |
leftmargin = 32 | |
rightmargin = 8 | |
svgwidth = len(heights) * (barwidth + gap) + 2 * leftmargin + rightmargin | |
svgheight = barheight + textheight | |
# create an SVG XML element | |
svg = et.Element('svg', width=str(svgwidth), height=str(svgheight), | |
version='1.1', xmlns='http://www.w3.org/2000/svg') | |
# Draw the bar graph | |
basey = svgheight - textheight | |
x = leftmargin | |
# Add units scale on left | |
if len(heights): | |
for h in [1, (max(heights) + 1) // 2, max(heights)]: | |
et.SubElement(svg, 'text', x='0', y='0', | |
style=('font-family:sans-serif;font-size:%dpx;' + | |
'text-anchor:end;alignment-baseline:hanging;' + | |
'transform:translate(%dpx, %dpx);') % | |
(textsize, x - gap, basey - h * unitheight)).text = str(h) | |
et.SubElement(svg, 'text', x='0', y='0', | |
style=('font-family:sans-serif;font-size:%dpx;' + | |
'text-anchor:middle;' + | |
'transform:translate(%dpx, %dpx) rotate(-90deg)') % | |
(textsize, x - gap - textsize, basey - h * unitheight / 2) | |
).text = 'units' | |
# Draw big category background rectangles | |
for catindex, (cat, catcount) in enumerate(categories): | |
if not catcount: | |
continue | |
et.SubElement(svg, 'rect', x=str(x), y=str(basey - rollup * unitheight), | |
width=(str((barwidth + gap) * catcount - gap)), | |
height = str(rollup*unitheight), | |
fill=bargraph_palette[catindex % len(bargraph_palette)][1]) | |
x += (barwidth + gap) * catcount | |
# Draw small bars as well as 45degree text labels | |
x = leftmargin | |
catindex = -1 | |
catcount = 0 | |
for label, height in zip(labels, heights): | |
while not catcount and catindex <= len(categories): | |
catindex += 1 | |
catcount = categories[catindex][1] | |
color = bargraph_palette[catindex % len(bargraph_palette)][0] | |
et.SubElement(svg, 'rect', x=str(x), y=str(basey-(height * unitheight)), | |
width=str(barwidth), height=str(height * unitheight), | |
fill=color) | |
x += barwidth | |
if show_labels: | |
et.SubElement(svg, 'text', x='0', y='0', | |
style=('font-family:sans-serif;font-size:%dpx;text-anchor:end;'+ | |
'transform:translate(%dpx, %dpx) rotate(-45deg);') % | |
(labelsize, x, basey + textmargin)).text = readable(label) | |
x += gap | |
catcount -= 1 | |
# Text labels for each category | |
x = leftmargin | |
for cat, catcount in categories: | |
if not catcount: | |
continue | |
et.SubElement(svg, 'text', x='0', y='0', | |
style=('font-family:sans-serif;font-size:%dpx;text-anchor:end;'+ | |
'transform:translate(%dpx, %dpx) rotate(-90deg);') % | |
(textsize, x + (barwidth + gap) * catcount - gap, | |
basey - rollup * unitheight + gap)).text = '%d %s' % ( | |
catcount, readable(cat + ('s' if catcount != 1 else ''))) | |
x += (barwidth + gap) * catcount | |
# Output - this is the bare svg. | |
result = et.tostring(svg) | |
if filename: | |
f = open(filename, 'wb') | |
# When writing to a file a special header is needed. | |
f.write(''.join([ | |
'<?xml version=\"1.0\" standalone=\"no\"?>\n', | |
'<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n', | |
'\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n'] | |
).encode('utf-8')) | |
f.write(result) | |
f.close() | |
return result | |
readable_replacements = [(re.compile(r[0]), r[1]) for r in [ | |
(r'-[sc]$', ''), | |
(r'_', ' '), | |
]] | |
def readable(label): | |
for pattern, subst in readable_replacements: | |
label= re.sub(pattern, subst, label) | |
return label | |
def reverse_normalize_from_transform(transform): | |
''' | |
Crawl around the transforms attached to a dataset looking for a | |
Normalize transform, and return it a corresponding ReverseNormalize, | |
or None if no normalization is found. | |
''' | |
if isinstance(transform, torchvision.transforms.Normalize): | |
return ReverseNormalize(transform.mean, transform.std) | |
t = getattr(transform, 'transform', None) | |
if t is not None: | |
return reverse_normalize_from_transform(t) | |
transforms = getattr(transform, 'transforms', None) | |
if transforms is not None: | |
for t in reversed(transforms): | |
result = reverse_normalize_from_transform(t) | |
if result is not None: | |
return result | |
return None | |
class ReverseNormalize: | |
''' | |
Applies the reverse of torchvision.transforms.Normalize. | |
''' | |
def __init__(self, mean, stdev): | |
mean = numpy.array(mean) | |
stdev = numpy.array(stdev) | |
self.mean = torch.from_numpy(mean)[None,:,None,None].float() | |
self.stdev = torch.from_numpy(stdev)[None,:,None,None].float() | |
def __call__(self, data): | |
device = data.device | |
return data.mul(self.stdev.to(device)).add_(self.mean.to(device)) | |
class ImageOnlySegRunner: | |
def __init__(self, dataset, recover_image=None): | |
if recover_image is None: | |
recover_image = reverse_normalize_from_transform(dataset) | |
self.recover_image = recover_image | |
self.dataset = dataset | |
def get_label_and_category_names(self): | |
return [('-', '-')], ['-'] | |
def run_and_segment_batch(self, batch, model, | |
want_bincount=False, want_rgb=False): | |
[im] = batch | |
device = next(model.parameters()).device | |
if want_rgb: | |
rgb = self.recover_image(im.clone() | |
).permute(0, 2, 3, 1).mul_(255).clamp(0, 255).byte() | |
else: | |
rgb = None | |
# Stubs for seg and bc | |
seg = torch.zeros(im.shape[0], 1, 1, 1, dtype=torch.long) | |
bc = torch.ones(im.shape[0], 1, dtype=torch.long) | |
# Run the model. | |
model(im.to(device)) | |
return seg, bc, rgb, im.shape[2:] | |
class ClassifierSegRunner: | |
def __init__(self, dataset, recover_image=None): | |
# The dataset contains explicit segmentations | |
if recover_image is None: | |
recover_image = reverse_normalize_from_transform(dataset) | |
self.recover_image = recover_image | |
self.dataset = dataset | |
def get_label_and_category_names(self): | |
catnames = self.dataset.categories | |
label_and_cat_names = [(readable(label), | |
catnames[self.dataset.label_category[i]]) | |
for i, label in enumerate(self.dataset.labels)] | |
return label_and_cat_names, catnames | |
def run_and_segment_batch(self, batch, model, | |
want_bincount=False, want_rgb=False): | |
''' | |
Runs the dissected model on one batch of the dataset, and | |
returns a multilabel semantic segmentation for the data. | |
Given a batch of size (n, c, y, x) the segmentation should | |
be a (long integer) tensor of size (n, d, y//r, x//r) where | |
d is the maximum number of simultaneous labels given to a pixel, | |
and where r is some (optional) resolution reduction factor. | |
In the segmentation returned, the label `0` is reserved for | |
the background "no-label". | |
In addition to the segmentation, bc, rgb, and shape are returned | |
where bc is a per-image bincount counting returned label pixels, | |
rgb is a viewable (n, y, x, rgb) byte image tensor for the data | |
for visualizations (reversing normalizations, for example), and | |
shape is the (y, x) size of the data. If want_bincount or | |
want_rgb are False, those return values may be None. | |
''' | |
im, seg, bc = batch | |
device = next(model.parameters()).device | |
if want_rgb: | |
rgb = self.recover_image(im.clone() | |
).permute(0, 2, 3, 1).mul_(255).clamp(0, 255).byte() | |
else: | |
rgb = None | |
# Run the model. | |
model(im.to(device)) | |
return seg, bc, rgb, im.shape[2:] | |
class GeneratorSegRunner: | |
def __init__(self, segmenter): | |
# The segmentations are given by an algorithm | |
if segmenter is None: | |
segmenter = UnifiedParsingSegmenter(segsizes=[256], segdiv='quad') | |
self.segmenter = segmenter | |
self.num_classes = len(segmenter.get_label_and_category_names()[0]) | |
def get_label_and_category_names(self): | |
return self.segmenter.get_label_and_category_names() | |
def run_and_segment_batch(self, batch, model, | |
want_bincount=False, want_rgb=False): | |
''' | |
Runs the dissected model on one batch of the dataset, and | |
returns a multilabel semantic segmentation for the data. | |
Given a batch of size (n, c, y, x) the segmentation should | |
be a (long integer) tensor of size (n, d, y//r, x//r) where | |
d is the maximum number of simultaneous labels given to a pixel, | |
and where r is some (optional) resolution reduction factor. | |
In the segmentation returned, the label `0` is reserved for | |
the background "no-label". | |
In addition to the segmentation, bc, rgb, and shape are returned | |
where bc is a per-image bincount counting returned label pixels, | |
rgb is a viewable (n, y, x, rgb) byte image tensor for the data | |
for visualizations (reversing normalizations, for example), and | |
shape is the (y, x) size of the data. If want_bincount or | |
want_rgb are False, those return values may be None. | |
''' | |
device = next(model.parameters()).device | |
z_batch = batch[0] | |
tensor_images = model(z_batch.to(device)) | |
seg = self.segmenter.segment_batch(tensor_images, downsample=2) | |
if want_bincount: | |
index = torch.arange(z_batch.shape[0], | |
dtype=torch.long, device=device) | |
bc = (seg + index[:, None, None, None] * self.num_classes).view(-1 | |
).bincount(minlength=z_batch.shape[0] * self.num_classes) | |
bc = bc.view(z_batch.shape[0], self.num_classes) | |
else: | |
bc = None | |
if want_rgb: | |
images = ((tensor_images + 1) / 2 * 255) | |
rgb = images.permute(0, 2, 3, 1).clamp(0, 255).byte() | |
else: | |
rgb = None | |
return seg, bc, rgb, tensor_images.shape[2:] | |