'''
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:]