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