def __init__(self, args, dissectdir=None, device=None):
self.cachedir = os.path.join(dissectdir, 'cache')
self.device = device if device is not None else torch.device('cpu')
self.dissectdir = dissectdir
self.modellock = threading.Lock()
# Load the generator from the pth file.
args_copy = EasyDict(args)
args_copy.edit = True
model = create_instrumented_model(args_copy)
model.eval()
self.model = model
# Get the set of layers of interest.
# Default: all shallow children except last.
self.layers = sorted(model.retained.keys())
# Move it to CUDA if wanted.
model.to(device)
self.quantiles = {
layer:
load_quantile_if_present(os.path.join(self.dissectdir,
safe_dir_name(layer)),
'quantiles.npz',
device=torch.device('cpu'))
for layer in self.layers
}
def add_ace_ranking_to_dissection(outdir, layer, classname, total_scores):
source_filename = os.path.join(outdir, 'dissect.json')
source_filename_bak = os.path.join(outdir, 'dissect.json.bak')
# Back up the dissection (if not already backed up) before modifying
if not os.path.exists(source_filename_bak):
shutil.copy(source_filename, source_filename_bak)
with open(source_filename) as f:
dissection = EasyDict(json.load(f))
ranking_name = '%s-ace' % classname
# Remove any old ace ranking with the same name
lrec = [l for l in dissection.layers if l.layer == layer][0]
lrec.rankings = [r for r in lrec.rankings if r.name != ranking_name]
# Now convert ace scores to rankings
new_rankings = [dict(
name=ranking_name,
score=(-total_scores).flatten().tolist(),
metric='ace')]
# Prepend to list.
lrec.rankings[2:2] = new_rankings
# Replace the old dissect.json in-place
with open(source_filename, 'w') as f:
json.dump(dissection, f, indent=1)
def load_projects(directory):
"""
searches for CONFIG_FILE_NAME in all subdirectories of directory
and creates data handlers for all of them
:param directory: scan directory
:return: null
"""
project_dirs = []
# Don't search more than 2 dirs deep.
search_depth = 2 + directory.count(os.path.sep)
for root, dirs, files in os.walk(directory):
if CONFIG_FILE_NAME in files:
project_dirs.append(root)
# Don't get subprojects under a project dir.
del dirs[:]
elif root.count(os.path.sep) >= search_depth:
del dirs[:]
for p_dir in project_dirs:
print('Loading %s' % os.path.join(p_dir, CONFIG_FILE_NAME))
with open(os.path.join(p_dir, CONFIG_FILE_NAME), 'r') as jf:
config = EasyDict(json.load(jf))
dh_id = os.path.split(p_dir)[1]
projects[dh_id] = DissectionProject(
config=config,
project_dir=p_dir,
path_url='data/' + os.path.relpath(p_dir, directory),
public_host=args.public_host)
def initial_ablation(args, dissectdir):
# Load initialization from dissection, based on iou scores.
with open(os.path.join(dissectdir, 'dissect.json')) as f:
dissection = EasyDict(json.load(f))
lrec = [l for l in dissection.layers if l.layer == args.layer][0]
rrec = [r for r in lrec.rankings if r.name == '%s-iou' % args.classname][0]
init_scores = -torch.tensor(rrec.score)
return init_scores / init_scores.max()
def load_segmentation_model(modeldir, segmodel_arch, segvocab, epoch=None):
# Load csv of class names
segmodel_dir = 'dataset/segmodel/%s-%s-%s' % ((segvocab, ) + segmodel_arch)
with open(os.path.join(segmodel_dir, 'labels.json')) as f:
labeldata = EasyDict(json.load(f))
# Automatically pick the last epoch available.
if epoch is None:
choices = [
os.path.basename(n)[14:-4] for n in glob.glob(
os.path.join(segmodel_dir, 'encoder_epoch_*.pth'))
]
epoch = max([int(c) for c in choices if c.isdigit()])
# Create a segmentation model
segbuilder = segmodel_module.ModelBuilder()
# example segmodel_arch = ('resnet101', 'upernet')
seg_encoder = segbuilder.build_encoder(arch=segmodel_arch[0],
fc_dim=2048,
weights=os.path.join(
segmodel_dir,
'encoder_epoch_%d.pth' % epoch))
seg_decoder = segbuilder.build_decoder(arch=segmodel_arch[1],
fc_dim=2048,
inference=True,
num_class=len(labeldata.labels),
weights=os.path.join(
segmodel_dir,
'decoder_epoch_%d.pth' % epoch))
segmodel = segmodel_module.SegmentationModule(
seg_encoder, seg_decoder, torch.nn.NLLLoss(ignore_index=-1))
segmodel.categories = [cat.name for cat in labeldata.categories]
segmodel.labels = [label.name for label in labeldata.labels]
categories = OrderedDict()
label_category = numpy.zeros(len(segmodel.labels), dtype=int)
for i, label in enumerate(labeldata.labels):
label_category[i] = segmodel.categories.index(label.category)
segmodel.meta = labeldata
segmodel.eval()
return segmodel
def create_instrumented_model(args, **kwargs):
'''
Creates an instrumented model out of a namespace of arguments that
correspond to ArgumentParser command-line args:
model: a string to evaluate as a constructor for the model.
pthfile: (optional) filename of .pth file for the model.
layers: a list of layers to instrument, defaulted if not provided.
edit: True to instrument the layers for editing.
gen: True for a generator model. One-pixel input assumed.
imgsize: For non-generator models, (y, x) dimensions for RGB input.
cuda: True to use CUDA.
The constructed model will be decorated with the following attributes:
input_shape: (usually 4d) tensor shape for single-image input.
output_shape: 4d tensor shape for output.
feature_shape: map of layer names to 4d tensor shape for featuremaps.
retained: map of layernames to tensors, filled after every evaluation.
ablation: if editing, map of layernames to [0..1] alpha values to fill.
replacement: if editing, map of layernames to values to fill.
When editing, the feature value x will be replaced by:
`x = (replacement * ablation) + (x * (1 - ablation))`
'''
args = EasyDict(vars(args), **kwargs)
# Construct the network
if args.model is None:
print_progress('No model specified')
return None
if isinstance(args.model, torch.nn.Module):
model = args.model
else:
model = autoimport_eval(args.model)
# Unwrap any DataParallel-wrapped model
if isinstance(model, torch.nn.DataParallel):
model = next(model.children())
# Load its state dict
meta = {}
if getattr(args, 'pthfile', None) is not None:
data = torch.load(args.pthfile)
if 'state_dict' in data:
meta = {}
for key in data:
if isinstance(data[key], numbers.Number):
meta[key] = data[key]
data = data['state_dict']
model.load_state_dict(data)
model.meta = meta
# Decide which layers to instrument.
if getattr(args, 'layer', None) is not None:
args.layers = [args.layer]
if getattr(args, 'layers', None) is None:
# Skip wrappers with only one named model
container = model
prefix = ''
while len(list(container.named_children())) == 1:
name, container = next(container.named_children())
prefix += name + '.'
# Default to all nontrivial top-level layers except last.
args.layers = [
prefix + name for name, module in container.named_children()
if type(module).__module__ not in [
# Skip ReLU and other activations.
'torch.nn.modules.activation',
# Skip pooling layers.
'torch.nn.modules.pooling'
]
][:-1]
print_progress('Defaulting to layers: %s' % ' '.join(args.layers))
# Instrument the layers.
retain_layers(model, args.layers)
if getattr(args, 'edit', False):
edit_layers(model, args.layers)
model.eval()
if args.cuda:
model.cuda()
# Annotate input, output, and feature shapes
annotate_model_shapes(model,
gen=getattr(args, 'gen', False),
imgsize=getattr(args, 'imgsize', None))
return model
def run_command(args):
verbose_progress(True)
progress = default_progress()
classname = args.classname # 'door'
layer = args.layer # 'layer4'
num_eval_units = 20
assert os.path.isfile(os.path.join(args.outdir, 'dissect.json')), (
"Should be a dissection directory")
if args.variant is None:
args.variant = 'ace'
if args.l2_lambda != 0.005:
args.variant = '%s_reg%g' % (args.variant, args.l2_lambda)
cachedir = os.path.join(args.outdir, safe_dir_name(layer), args.variant,
classname)
if pidfile_taken(os.path.join(cachedir, 'lock.pid'), True):
sys.exit(0)
# Take defaults for model constructor etc from dissect.json settings.
with open(os.path.join(args.outdir, 'dissect.json')) as f:
dissection = EasyDict(json.load(f))
if args.model is None:
args.model = dissection.settings.model
if args.pthfile is None:
args.pthfile = dissection.settings.pthfile
if args.segmenter is None:
args.segmenter = dissection.settings.segmenter
# Default segmenter class
if args.segmenter is None:
args.segmenter = ("netdissect.segmenter.UnifiedParsingSegmenter(" +
"segsizes=[256], segdiv='quad')")
if (not args.no_cache and
os.path.isfile(os.path.join(cachedir, 'snapshots', 'epoch-%d.npy' % (
args.train_epochs - 1))) and
os.path.isfile(os.path.join(cachedir, 'report.json'))):
print('%s already done' % cachedir)
sys.exit(0)
os.makedirs(cachedir, exist_ok=True)
# Instantiate generator
model = create_instrumented_model(args, gen=True, edit=True,
layers=[args.layer])
if model is None:
print('No model specified')
sys.exit(1)
# Instantiate segmenter
segmenter = autoimport_eval(args.segmenter)
labelnames, catname = segmenter.get_label_and_category_names()
classnum = [i for i, (n, c) in enumerate(labelnames) if n == classname][0]
num_classes = len(labelnames)
with open(os.path.join(cachedir, 'labelnames.json'), 'w') as f:
json.dump(labelnames, f, indent=1)
# Sample sets for training.
full_sample = netdissect.zdataset.z_sample_for_model(model,
args.search_size, seed=10)
second_sample = netdissect.zdataset.z_sample_for_model(model,
args.search_size, seed=11)
# Load any cached data.
cache_filename = os.path.join(cachedir, 'corpus.npz')
corpus = EasyDict()
try:
if not args.no_cache:
corpus = EasyDict({k: torch.from_numpy(v)
for k, v in numpy.load(cache_filename).items()})
except:
pass
# The steps for the computation.
compute_present_locations(args, corpus, cache_filename,
model, segmenter, classnum, full_sample)
compute_mean_present_features(args, corpus, cache_filename, model)
compute_feature_quantiles(args, corpus, cache_filename, model, full_sample)
compute_candidate_locations(args, corpus, cache_filename, model, segmenter,
classnum, second_sample)
# visualize_training_locations(args, corpus, cachedir, model)
init_ablation = initial_ablation(args, args.outdir)
scores = train_ablation(args, corpus, cache_filename,
model, segmenter, classnum, init_ablation)
summarize_scores(args, corpus, cachedir, layer, classname,
args.variant, scores)
if args.variant == 'ace':
add_ace_ranking_to_dissection(args.outdir, layer, classname, scores)
def optimize_neurons(self):
# Set up console output
verbose_progress(True)
gan_model = self.generator.model
annotate_model_shapes(gan_model, gen=True)
outdir = os.path.join(
self.args.results, 'dissect',
self.args.name_checkpoint + '_' + str(time.time()))
os.makedirs(outdir, exist_ok=True)
size = 1000
sample = z_sample_for_model(gan_model, size)
train_sample = z_sample_for_model(gan_model, size, seed=2)
dataset = TensorDataset(sample)
train_dataset = TensorDataset(train_sample)
self.cluster_segmenter = ClusterSegmenter(self.model, self.clusters,
self.mean_clust,
self.std_clust)
segrunner = GeneratorSegRunner(self.cluster_segmenter)
netname = outdir
# Run dissect
with torch.no_grad():
dissect(
outdir,
gan_model,
dataset,
train_dataset=train_dataset,
segrunner=segrunner,
examples_per_unit=20,
netname=netname,
quantile_threshold='iqr',
meta=None,
make_images=False, # True,
make_labels=True,
make_maxiou=False,
make_covariance=False,
make_report=True,
make_row_images=True,
make_single_images=True,
batch_size=8,
num_workers=8,
rank_all_labels=True)
sample_ablate = z_sample_for_model(gan_model, 16)
dataset_ablate = TensorDataset(sample_ablate)
data_loader = torch.utils.data.DataLoader(dataset_ablate,
batch_size=8,
shuffle=False,
num_workers=8,
pin_memory=True,
sampler=None)
with open(os.path.join(outdir, 'dissect.json')) as f:
data = EasyDict(json.load(f))
dissect_layer = {lrec.layer: lrec for lrec in data.layers}
self.layers_units = {
'layer2': [],
'layer3': [],
'layer4': [],
'layer5': [],
'layer6': [],
}
noise_units = np.array([35, 221, 496, 280])
for i in range(2, len(self.clusters) + 2):
print('Cluster', i)
rank_name = 'c_{0}-iou'.format(i)
for l in range(len(self.layer_list_all)):
ranking = next(
r
for r in dissect_layer[self.layer_list_all[l]].rankings
if r.name == rank_name)
unit_list = np.array(range(512))
unit_list[noise_units] = 0
ordering = np.argsort(ranking.score)
units_list = unit_list[ordering]
self.layers_units[self.layer_list_all[l]].append(
units_list)
# Mark the directory so that it's not done again.
mark_job_done(outdir)
def main():
# Training settings
def strpair(arg):
p = tuple(arg.split(':'))
if len(p) == 1:
p = p + p
return p
parser = argparse.ArgumentParser(
description='Ablation eval',
epilog=textwrap.dedent(help_epilog),
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--model',
type=str,
default=None,
help='constructor for the model to test')
parser.add_argument('--pthfile',
type=str,
default=None,
help='filename of .pth file for the model')
parser.add_argument('--outdir',
type=str,
default='dissect',
required=True,
help='directory for dissection output')
parser.add_argument('--layers',
type=strpair,
nargs='+',
help='space-separated list of layer names to edit' +
', in the form layername[:reportedname]')
parser.add_argument('--classes',
type=str,
nargs='+',
help='space-separated list of class names to ablate')
parser.add_argument('--metric',
type=str,
default='iou',
help='ordering metric for selecting units')
parser.add_argument('--unitcount',
type=int,
default=30,
help='number of units to ablate')
parser.add_argument('--segmenter',
type=str,
help='directory containing segmentation dataset')
parser.add_argument('--netname',
type=str,
default=None,
help='name for network in generated reports')
parser.add_argument('--batch_size',
type=int,
default=5,
help='batch size for forward pass')
parser.add_argument('--size',
type=int,
default=200,
help='number of images to test')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA usage')
parser.add_argument('--quiet',
action='store_true',
default=False,
help='silences console output')
if len(sys.argv) == 1:
parser.print_usage(sys.stderr)
sys.exit(1)
args = parser.parse_args()
# Set up console output
pbar.verbose(not args.quiet)
# Speed up pytorch
torch.backends.cudnn.benchmark = True
# Set up CUDA
args.cuda = not args.no_cuda and torch.cuda.is_available()
if args.cuda:
torch.backends.cudnn.benchmark = True
# Take defaults for model constructor etc from dissect.json settings.
with open(os.path.join(args.outdir, 'dissect.json')) as f:
dissection = EasyDict(json.load(f))
if args.model is None:
args.model = dissection.settings.model
if args.pthfile is None:
args.pthfile = dissection.settings.pthfile
if args.segmenter is None:
args.segmenter = dissection.settings.segmenter
# Instantiate generator
model = create_instrumented_model(args, gen=True, edit=True)
if model is None:
print('No model specified')
sys.exit(1)
# Instantiate model
device = next(model.parameters()).device
input_shape = model.input_shape
# 4d input if convolutional, 2d input if first layer is linear.
raw_sample = standard_z_sample(args.size, input_shape[1],
seed=2).view((args.size, ) +
input_shape[1:])
dataset = TensorDataset(raw_sample)
# Create the segmenter
segmenter = autoimport_eval(args.segmenter)
# Now do the actual work.
labelnames, catnames = (segmenter.get_label_and_category_names(dataset))
label_category = [
catnames.index(c) if c in catnames else 0 for l, c in labelnames
]
labelnum_from_name = {n[0]: i for i, n in enumerate(labelnames)}
segloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
num_workers=10,
pin_memory=(device.type == 'cuda'))
# Index the dissection layers by layer name.
dissect_layer = {lrec.layer: lrec for lrec in dissection.layers}
# First, collect a baseline
for l in model.ablation:
model.ablation[l] = None
# For each sort-order, do an ablation
for classname in pbar(args.classes):
pbar.post(c=classname)
for layername in pbar(model.ablation):
pbar.post(l=layername)
rankname = '%s-%s' % (classname, args.metric)
classnum = labelnum_from_name[classname]
try:
ranking = next(r for r in dissect_layer[layername].rankings
if r.name == rankname)
except:
print('%s not found' % rankname)
sys.exit(1)
ordering = numpy.argsort(ranking.score)
# Check if already done
ablationdir = os.path.join(args.outdir, layername, 'pixablation')
if os.path.isfile(os.path.join(ablationdir, '%s.json' % rankname)):
with open(os.path.join(ablationdir,
'%s.json' % rankname)) as f:
data = EasyDict(json.load(f))
# If the unit ordering is not the same, something is wrong
if not all(a == o
for a, o in zip(data.ablation_units, ordering)):
continue
if len(data.ablation_effects) >= args.unitcount:
continue # file already done.
measurements = data.ablation_effects
measurements = measure_ablation(segmenter, segloader, model,
classnum, layername,
ordering[:args.unitcount])
measurements = measurements.cpu().numpy().tolist()
os.makedirs(ablationdir, exist_ok=True)
with open(os.path.join(ablationdir, '%s.json' % rankname),
'w') as f:
json.dump(
dict(classname=classname,
classnum=classnum,
baseline=measurements[0],
layer=layername,
metric=args.metric,
ablation_units=ordering.tolist(),
ablation_effects=measurements[1:]), f)
for scene, classname in [
('churchoutdoor', 'door'),
('churchoutdoor', 'tree'),
]:
fig = Figure(figsize=(4.5, 3.5))
FigureCanvas(fig)
ax = fig.add_subplot(111)
for metric in [test_metric, 'iou']:
# Plot line graph
with open(
os.path.join(basedir, scene, layername, 'ablation',
'%s-%s.json' % (classname, metric))) as f:
data = EasyDict(json.load(f))
ax.plot(
[0] +
[(data.baseline - data.ablation_effects[str(i)]) / data.baseline
for i in range(1, xlim + 1)],
label='Top units by IoU' if metric is 'iou' else 'Units by ACE')
ax.set_title('Effect of ablating units for %s' % (classname))
ax.xaxis.set_major_locator(MaxNLocator(integer=True))
ax.grid(True)
ax.legend()
ax.set_ylim(0, 1.0)
ax.set_xlim(0, xlim)
ax.set_xticks(range(0, xlim + 1, 5))
ax.set_ylabel('Portion of %s pixels removed' % classname)
ax.set_xlabel('Number of units ablated')
def main():
# Training settings
def strpair(arg):
p = tuple(arg.split(':'))
if len(p) == 1:
p = p + p
return p
parser = argparse.ArgumentParser(description='Ablation eval',
epilog=textwrap.dedent(help_epilog),
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--model', type=str, default=None,
help='constructor for the model to test')
parser.add_argument('--pthfile', type=str, default=None,
help='filename of .pth file for the model')
parser.add_argument('--outdir', type=str, default='dissect', required=True,
help='directory for dissection output')
parser.add_argument('--layer', type=strpair,
help='space-separated list of layer names to edit' +
', in the form layername[:reportedname]')
parser.add_argument('--classname', type=str,
help='class name to ablate')
parser.add_argument('--metric', type=str, default='iou',
help='ordering metric for selecting units')
parser.add_argument('--unitcount', type=int, default=30,
help='number of units to ablate')
parser.add_argument('--segmenter', type=str,
help='directory containing segmentation dataset')
parser.add_argument('--netname', type=str, default=None,
help='name for network in generated reports')
parser.add_argument('--batch_size', type=int, default=25,
help='batch size for forward pass')
parser.add_argument('--mixed_units', action='store_true', default=False,
help='true to keep alpha for non-zeroed units')
parser.add_argument('--size', type=int, default=200,
help='number of images to test')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA usage')
parser.add_argument('--quiet', action='store_true', default=False,
help='silences console output')
if len(sys.argv) == 1:
parser.print_usage(sys.stderr)
sys.exit(1)
args = parser.parse_args()
# Set up console output
verbose_progress(not args.quiet)
# Speed up pytorch
torch.backends.cudnn.benchmark = True
# Set up CUDA
args.cuda = not args.no_cuda and torch.cuda.is_available()
if args.cuda:
torch.backends.cudnn.benchmark = True
# Take defaults for model constructor etc from dissect.json settings.
with open(os.path.join(args.outdir, 'dissect.json')) as f:
dissection = EasyDict(json.load(f))
if args.model is None:
args.model = dissection.settings.model
if args.pthfile is None:
args.pthfile = dissection.settings.pthfile
if args.segmenter is None:
args.segmenter = dissection.settings.segmenter
if args.layer is None:
args.layer = dissection.settings.layers[0]
args.layers = [args.layer]
# Also load specific analysis
layername = args.layer[1]
if args.metric == 'iou':
summary = dissection
else:
with open(os.path.join(args.outdir, layername, args.metric,
args.classname, 'summary.json')) as f:
summary = EasyDict(json.load(f))
# Instantiate generator
model = create_instrumented_model(args, gen=True, edit=True)
if model is None:
print('No model specified')
sys.exit(1)
# Instantiate model
device = next(model.parameters()).device
input_shape = model.input_shape
# 4d input if convolutional, 2d input if first layer is linear.
raw_sample = standard_z_sample(args.size, input_shape[1], seed=3).view(
(args.size,) + input_shape[1:])
dataset = TensorDataset(raw_sample)
# Create the segmenter
segmenter = autoimport_eval(args.segmenter)
# Now do the actual work.
labelnames, catnames = (
segmenter.get_label_and_category_names(dataset))
label_category = [catnames.index(c) if c in catnames else 0
for l, c in labelnames]
labelnum_from_name = {n[0]: i for i, n in enumerate(labelnames)}
segloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size, num_workers=10,
pin_memory=(device.type == 'cuda'))
# Index the dissection layers by layer name.
# First, collect a baseline
for l in model.ablation:
model.ablation[l] = None
# For each sort-order, do an ablation
progress = default_progress()
classname = args.classname
classnum = labelnum_from_name[classname]
# Get iou ranking from dissect.json
iou_rankname = '%s-%s' % (classname, 'iou')
dissect_layer = {lrec.layer: lrec for lrec in dissection.layers}
iou_ranking = next(r for r in dissect_layer[layername].rankings
if r.name == iou_rankname)
# Get trained ranking from summary.json
rankname = '%s-%s' % (classname, args.metric)
summary_layer = {lrec.layer: lrec for lrec in summary.layers}
ranking = next(r for r in summary_layer[layername].rankings
if r.name == rankname)
# Get ordering, first by ranking, then break ties by iou.
ordering = [t[2] for t in sorted([(s1, s2, i)
for i, (s1, s2) in enumerate(zip(ranking.score, iou_ranking.score))])]
values = (-numpy.array(ranking.score))[ordering]
if not args.mixed_units:
values[...] = 1
ablationdir = os.path.join(args.outdir, layername, 'fullablation')
measurements = measure_full_ablation(segmenter, segloader,
model, classnum, layername,
ordering[:args.unitcount], values[:args.unitcount])
measurements = measurements.cpu().numpy().tolist()
os.makedirs(ablationdir, exist_ok=True)
with open(os.path.join(ablationdir, '%s.json'%rankname), 'w') as f:
json.dump(dict(
classname=classname,
classnum=classnum,
baseline=measurements[0],
layer=layername,
metric=args.metric,
ablation_units=ordering,
ablation_values=values.tolist(),
ablation_effects=measurements[1:]), f)
def create_instrumented_model(args, **kwargs):
'''
Creates an instrumented model out of a namespace of arguments that
correspond to ArgumentParser command-line args:
model: a string to evaluate as a constructor for the model.
pthfile: (optional) filename of .pth file for the model.
layers: a list of layers to instrument, defaulted if not provided.
edit: True to instrument the layers for editing.
gen: True for a generator model. One-pixel input assumed.
imgsize: For non-generator models, (y, x) dimensions for RGB input.
cuda: True to use CUDA.
The constructed model will be decorated with the following attributes:
input_shape: (usually 4d) tensor shape for single-image input.
output_shape: 4d tensor shape for output.
feature_shape: map of layer names to 4d tensor shape for featuremaps.
retained: map of layernames to tensors, filled after every evaluation.
ablation: if editing, map of layernames to [0..1] alpha values to fill.
replacement: if editing, map of layernames to values to fill.
When editing, the feature value x will be replaced by:
`x = (replacement * ablation) + (x * (1 - ablation))`
'''
args = EasyDict(vars(args), **kwargs)
# Construct the network
if args.model is None:
print_progress('No model specified')
return None
#if isinstance(args.model, torch.nn.Module):
# model = args.model
#else:
# model = autoimport_eval(args.model)
model = UnetNormalized()
checkpoint = torch.load('p2p_churches.pth')
model.load_state_dict(checkpoint)
model.cuda()
#model.eval()
torch.no_grad()
# Unwrap any DataParallel-wrapped model
if isinstance(model, torch.nn.DataParallel):
model = next(model.children())
# Load its state dict
#meta = {}
#if getattr(args, 'pthfile', None) is not None:
# data = torch.load(args.pthfile)
# if 'state_dict' in data:
# meta = {}
# for key in data:
# if isinstance(data[key], numbers.Number):
# meta[key] = data[key]
# data = data['state_dict']
# submodule = getattr(args, 'submodule', None)
# if submodule is not None and len(submodule):
# remove_prefix = submodule + '.'
# data = { k[len(remove_prefix):]: v for k, v in data.items()
# if k.startswith(remove_prefix)}
# if not len(data):
# print_progress('No submodule %s found in %s' %
# (submodule, args.pthfile))
# return None
# model.load_state_dict(data, strict=not getattr(args, 'unstrict', False))
# Decide which layers to instrument.
#if getattr(args, 'layer', None) is not None:
# args.layers = [args.layer]
args.layers = [layer5, layer9, layer12]
#if getattr(args, 'layers', None) is None:
# # Skip wrappers with only one named model
# container = model
# prefix = ''
# while len(list(container.named_children())) == 1:
# name, container = next(container.named_children())
# prefix += name + '.'
# # Default to all nontrivial top-level layers except last.
# args.layers = [prefix + name
# for name, module in container.named_children()
# if type(module).__module__ not in [
# # Skip ReLU and other activations.
# 'torch.nn.modules.activation',
# # Skip pooling layers.
# 'torch.nn.modules.pooling']
# ][:-1]
# print_progress('Defaulting to layers: %s' % ' '.join(args.layers))
# Now wrap the model for instrumentation.
model = InstrumentedModel(model)
# model.meta = meta
# Instrument the layers.
model.retain_layers(args.layers)
#model.eval()
#if args.cuda:
# model.cuda()
# Annotate input, output, and feature shapes
annotate_model_shapes(model,
gen=getattr(args, 'gen', False),
imgsize=getattr(args, 'imgsize', None))
return model
def main():
args = parseargs()
threshold_iou = args.miniou
layer_report = {}
qdir = '-%d' % (args.quantile * 1000) if args.quantile != 0.01 else ''
for layer in args.layers:
input_filename = 'old_results/%s-%s-%s-%s/report.json' % (
args.model, args.dataset, args.seg,
layer.replace('features.',''))
with open(input_filename) as f:
layer_report[layer] = EasyDict(json.load(f))
layer_scores = {}
for layer in args.layers:
dirname = 'results/%s/%s/%s/%s/%d/' % (
args.model, args.dataset, layer, args.seg,
args.samplesize)
layer_scores[layer] = load_intersect_iou(dirname)
# Now assemble the data needed for the graph
# (Layername, [(catname, [unitcount, unitcount, unitcount]), (catname..)
cat_order = ['object', 'part', 'material', 'color']
graph_data = []
for layer in args.layers:
layer_data = []
catmap = defaultdict(lambda: defaultdict(int))
units = layer_report[layer].get('units',
layer_report[layer].get('images', None)) # old format
for unitrec in units:
if unitrec.iou is None or unitrec.iou < threshold_iou:
continue
if layer_scores[layer][unitrec.unit] < 0.1:
continue
catmap[unitrec.cat][unitrec.label] += 1
for cat in cat_order:
if cat not in catmap:
continue
# For this graph we do not need labels
cat_data = list(catmap[cat].values())
cat_data.sort(key=lambda x: -x)
layer_data.append((cat, cat_data))
graph_data.append((layer.replace('features.', ''), layer_data))
# Now make the actual graph
largest_layer = max(sum(len(cat_data)
for cat, cat_data in layer_data)
for layer, layer_data in graph_data)
layer_height = 14
layer_gap = 2
barwidth = 3
bargap = 0
leftmargin = 48
margin = 8
svgwidth = largest_layer * (barwidth + bargap) + margin + leftmargin
svgheight = ((layer_height + layer_gap) * len(args.layers) - layer_gap +
2 * margin)
textsize = 10
# 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 big category background rectangles
y = margin
for layer, layer_data in graph_data:
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, leftmargin - 4, y + (layer_height - textsize) / 2)
).text = str(layer)
barmax = max(max(cat_data) if len(cat_data) else 1
for cat, cat_data in layer_data) if len(layer_data) else 1
barscale = float(layer_height) / barmax
x = leftmargin
for cat, cat_data in layer_data:
catwidth = len(cat_data) * (barwidth + bargap)
et.SubElement(svg, 'rect',
x=str(x), y=str(y),
width=str(catwidth),
height=str(layer_height),
fill=cat_palette[cat][1])
for bar in cat_data:
barheight = barscale * bar
et.SubElement(svg, 'rect',
x=str(x), y=str(y + layer_height - barheight),
width=str(barwidth),
height=str(barheight),
fill=cat_palette[cat][0])
x += barwidth + bargap
y += layer_height + layer_gap
# Output - this is the bare svg.
result = et.tostring(svg).decode('utf-8')
# Now add the file header.
result = ''.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',
result])
output_filename = 'results/multilayer-a.svg'
os.makedirs(os.path.dirname(output_filename), exist_ok=True)
print('writing to %s' % output_filename)
with open(output_filename, 'w') as f:
f.write(result)
def main():
# Training settings
def strpair(arg):
p = tuple(arg.split(':'))
if len(p) == 1:
p = p + p
return p
parser = argparse.ArgumentParser(
description='Net dissect utility',
epilog=textwrap.dedent(help_epilog),
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--model',
type=str,
default=None,
help='constructor for the model to test')
parser.add_argument('--pthfile',
type=str,
default=None,
help='filename of .pth file for the model')
parser.add_argument('--outdir',
type=str,
default='dissect',
help='directory for dissection output')
parser.add_argument('--layers',
type=strpair,
nargs='+',
help='space-separated list of layer names to edit' +
', in the form layername[:reportedname]')
parser.add_argument('--classes',
type=str,
nargs='+',
help='space-separated list of class names to ablate')
parser.add_argument('--metric',
type=str,
default='iou',
help='ordering metric for selecting units')
parser.add_argument('--startcount',
type=int,
default=1,
help='number of units to ablate')
parser.add_argument('--unitcount',
type=int,
default=30,
help='number of units to ablate')
parser.add_argument('--segmenter',
type=str,
default='dataset/broden',
help='directory containing segmentation dataset')
parser.add_argument('--netname',
type=str,
default=None,
help='name for network in generated reports')
parser.add_argument('--batch_size',
type=int,
default=5,
help='batch size for forward pass')
parser.add_argument('--size',
type=int,
default=1000,
help='number of images to test')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA usage')
parser.add_argument('--quiet',
action='store_true',
default=False,
help='silences console output')
if len(sys.argv) == 1:
parser.print_usage(sys.stderr)
sys.exit(1)
args = parser.parse_args()
# Set up console output
verbose_progress(not args.quiet)
# Speed up pytorch
torch.backends.cudnn.benchmark = True
# Construct the network
if args.model is None:
print_progress('No model specified')
sys.exit(1)
# Set up CUDA
args.cuda = not args.no_cuda and torch.cuda.is_available()
if args.cuda:
torch.backends.cudnn.benchmark = True
model = autoimport_eval(args.model)
# Unwrap any DataParallel-wrapped model
if isinstance(model, torch.nn.DataParallel):
model = next(model.children())
# Load its state dict
meta = {}
if args.pthfile is None:
print_progress('Dissecting model without pth file.')
else:
data = torch.load(args.pthfile)
if 'state_dict' in data:
meta = {}
for key in data:
if isinstance(data[key], numbers.Number):
meta[key] = data[key]
data = data['state_dict']
model.load_state_dict(data)
# Instrument it and prepare it for eval
if not args.layers:
# Skip wrappers with only one named modele
container = model
prefix = ''
while len(list(container.named_children())) == 1:
name, container = next(container.named_children())
prefix += name + '.'
# Default to all nontrivial top-level layers except last.
args.layers = [
prefix + name for name, module in container.named_children()
if type(module).__module__ not in [
# Skip ReLU and other activations.
'torch.nn.modules.activation',
# Skip pooling layers.
'torch.nn.modules.pooling'
]
][:-1]
print_progress('Defaulting to layers: %s' % ' '.join(args.layers))
edit_layers(model, args.layers)
model.eval()
if args.cuda:
model.cuda()
# Set up the output directory, verify write access
if args.outdir is None:
args.outdir = os.path.join('dissect', type(model).__name__)
print_progress('Writing output into %s.' % args.outdir)
os.makedirs(args.outdir, exist_ok=True)
train_dataset = None
# Examine first conv in model to determine input feature size.
first_layer = [
c for c in model.modules()
if isinstance(c, (torch.nn.Conv2d, torch.nn.ConvTranspose2d,
torch.nn.Linear))
][0]
# 4d input if convolutional, 2d input if first layer is linear.
if isinstance(first_layer, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)):
sample = standard_z_sample(args.size,
first_layer.in_channels)[:, :, None, None]
train_sample = standard_z_sample(args.size,
first_layer.in_channels,
seed=2)[:, :, None, None]
else:
sample = standard_z_sample(args.size, first_layer.in_features)
train_sample = standard_z_sample(args.size,
first_layer.in_features,
seed=2)
dataset = TensorDataset(sample)
train_dataset = TensorDataset(train_sample)
recovery = autoimport_eval(args.segmenter)
# Now do the actual work.
device = next(model.parameters()).device
labelnames, catnames = (recovery.get_label_and_category_names(dataset))
label_category = [catnames.index(c) for l, c in labelnames]
labelnum_from_name = {n[0]: i for i, n in enumerate(labelnames)}
segloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
num_workers=10,
pin_memory=(device.type == 'cuda'))
with open(os.path.join(args.outdir, 'dissect.json'), 'r') as f:
dissect = EasyDict(json.load(f))
# Index the dissection layers by layer name.
dissect_layer = {lrec.layer: lrec for lrec in dissect.layers}
# First, collect a baseline
for l in model.ablation:
model.ablation[l] = None
baseline = count_segments(recovery, segloader, model)
# For each sort-order, do an ablation
progress = default_progress()
for classname in progress(args.classes):
post_progress(c=classname)
for layername in progress(model.ablation):
post_progress(l=layername)
rankname = '%s-%s' % (classname, args.metric)
measurements = {}
classnum = labelnum_from_name[classname]
try:
ranking = next(r for r in dissect_layer[layername].rankings
if r.name == rankname)
except:
print('%s not found' % rankname)
sys.exit(1)
ordering = numpy.argsort(ranking.score)
# Check if already done
ablationdir = os.path.join(args.outdir, layername, 'ablation')
if os.path.isfile(os.path.join(ablationdir, '%s.json' % rankname)):
with open(os.path.join(ablationdir,
'%s.json' % rankname)) as f:
data = EasyDict(json.load(f))
# If the unit ordering is not the same, something is wrong
if not all(a == o
for a, o in zip(data.ablation_units, ordering)):
import pdb
pdb.set_trace()
continue
if len(data.ablation_effects) >= args.unitcount:
continue # file already done.
measurements = data.ablation_effects
for count in progress(range(args.startcount,
min(args.unitcount, len(ordering)) +
1),
desc='units'):
if str(count) in measurements:
continue
ablation = numpy.zeros(len(ranking.score), dtype='float32')
ablation[ordering[:count]] = 1
for l in model.ablation:
model.ablation[l] = ablation if layername == l else None
m = count_segments(recovery, segloader, model)[classnum].item()
print_progress(
'%s %s %d units (#%d), %g -> %g' %
(layername, rankname, count, ordering[count - 1].item(),
baseline[classnum].item(), m))
measurements[str(count)] = m
os.makedirs(ablationdir, exist_ok=True)
with open(os.path.join(ablationdir, '%s.json' % rankname),
'w') as f:
json.dump(
dict(classname=classname,
classnum=classnum,
baseline=baseline[classnum].item(),
layer=layername,
metric=args.metric,
ablation_units=ordering.tolist(),
ablation_effects=measurements), f)