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)
from torchvision import transforms
from netdissect.nethook import InstrumentedModel
from netdissect.autoeval import autoimport_eval
from netdissect.modelconfig import annotate_model_shapes
from netdissect.segviz import segment_visualization
import PIL
batch_size = 1
data = get_segments_dataset('dataset/Adissect_toy')
#model = autoimport_eval("p2pgan.from_pth_file('models/pix2pix/p2p_churches.pth')")
model = from_pth_file('models/pix2pix/p2p_churches.pth')
segmenter = (
"netdissect.segmenter.UnifiedParsingSegmenter(segsizes=[256], segdiv='quad')"
)
segrunner = GeneratorSegRunner(autoimport_eval(segmenter))
layer5 = ('model.model.1.model.3.model.3.model.3.model.1', 'layer5')
layer9 = (
'model.model.1.model.3.model.3.model.3.model.3.model.3.model.3.model.3',
'layer9')
layer12 = ('model.model.1.model.3.model.3.model.3.model.5', 'layer12')
model = InstrumentedModel(model)
model.retain_layers([layer5, layer9, layer12])
annotate_model_shapes(model, gen=True, imgsize=None)
segloader = torch.utils.data.DataLoader(data,
batch_size=batch_size,
pin_memory=True)
def main():
# Training settings
def strpair(arg):
p = tuple(arg.split(':'))
if len(p) == 1:
p = p + p
return p
def intpair(arg):
p = arg.split(',')
if len(p) == 1:
p = p + p
return tuple(int(v) for v in p)
parser = argparse.ArgumentParser(
description='Net dissect utility',
prog='python -m netdissect',
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('--unstrict',
action='store_true',
default=False,
help='ignore unexpected pth parameters')
parser.add_argument('--submodule',
type=str,
default=None,
help='submodule to load from pthfile')
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 dissect' +
', in the form layername[:reportedname]')
parser.add_argument('--segments',
type=str,
default='dataset/broden',
help='directory containing segmentation dataset')
parser.add_argument('--segmenter',
type=str,
default=None,
help='constructor for asegmenter class')
parser.add_argument('--download',
action='store_true',
default=False,
help='downloads Broden dataset if needed')
parser.add_argument('--imagedir',
type=str,
default=None,
help='directory containing image-only dataset')
parser.add_argument('--imgsize',
type=intpair,
default=(227, 227),
help='input image size to use')
parser.add_argument('--netname',
type=str,
default=None,
help='name for network in generated reports')
parser.add_argument('--meta',
type=str,
nargs='+',
help='json files of metadata to add to report')
parser.add_argument('--merge',
type=str,
help='json file of unit data to merge in report')
parser.add_argument('--examples',
type=int,
default=20,
help='number of image examples per unit')
parser.add_argument('--size',
type=int,
default=10000,
help='dataset subset size to use')
parser.add_argument('--batch_size',
type=int,
default=100,
help='batch size for forward pass')
parser.add_argument('--num_workers',
type=int,
default=24,
help='number of DataLoader workers')
parser.add_argument('--quantile_threshold',
type=strfloat,
default=None,
choices=[FloatRange(0.0, 1.0), 'iqr'],
help='quantile to use for masks')
parser.add_argument('--no-labels',
action='store_true',
default=False,
help='disables labeling of units')
parser.add_argument('--maxiou',
action='store_true',
default=False,
help='enables maxiou calculation')
parser.add_argument('--covariance',
action='store_true',
default=False,
help='enables covariance calculation')
parser.add_argument('--rank_all_labels',
action='store_true',
default=False,
help='include low-information labels in rankings')
parser.add_argument('--no-images',
action='store_true',
default=False,
help='disables generation of unit images')
parser.add_argument('--no-report',
action='store_true',
default=False,
help='disables generation report summary')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA usage')
parser.add_argument('--gen',
action='store_true',
default=False,
help='test a generator model (e.g., a GAN)')
parser.add_argument('--gan',
action='store_true',
default=False,
help='synonym for --gen')
parser.add_argument('--perturbation',
default=None,
help='filename of perturbation attack to apply')
parser.add_argument('--add_scale_offset',
action='store_true',
default=None,
help='offsets masks according to stride and padding')
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()
args.images = not args.no_images
args.report = not args.no_report
args.labels = not args.no_labels
if args.gan:
args.gen = args.gan
# Set up console output
verbose_progress(not args.quiet)
# Exit right away if job is already done or being done.
if args.outdir is not None:
exit_if_job_done(args.outdir)
# Speed up pytorch
torch.backends.cudnn.benchmark = True
# Special case: download flag without model to test.
if args.model is None and args.download:
from netdissect.broden import ensure_broden_downloaded
for resolution in [224, 227, 384]:
ensure_broden_downloaded(args.segments, resolution, 1)
from netdissect.segmenter import ensure_upp_segmenter_downloaded
ensure_upp_segmenter_downloaded('dataset/segmodel')
sys.exit(0)
# Help if broden is not present
if not args.gen and not args.imagedir and not os.path.isdir(args.segments):
print_progress('Segmentation dataset not found at %s.' % args.segments)
print_progress('Specify dataset directory using --segments [DIR]')
print_progress('To download Broden, run: netdissect --download')
sys.exit(1)
# Default segmenter class
if args.gen and args.segmenter is None:
args.segmenter = ("netdissect.segmenter.UnifiedParsingSegmenter(" +
"segsizes=[256], segdiv=None)")
# Default threshold
if args.quantile_threshold is None:
if args.gen:
args.quantile_threshold = 'iqr'
else:
args.quantile_threshold = 0.005
# Set up CUDA
args.cuda = not args.no_cuda and torch.cuda.is_available()
if args.cuda:
torch.backends.cudnn.benchmark = True
# Construct the network with specified layers instrumented
if args.model is None:
print_progress('No model specified')
sys.exit(1)
model = create_instrumented_model(args)
# Update any metadata from files, if any
meta = getattr(model, 'meta', {})
if args.meta:
for mfilename in args.meta:
with open(mfilename) as f:
meta.update(json.load(f))
# Load any merge data from files
mergedata = None
if args.merge:
with open(args.merge) as f:
mergedata = json.load(f)
# Set up the output directory, verify write access
if args.outdir is None:
args.outdir = os.path.join('dissect', type(model).__name__)
exit_if_job_done(args.outdir)
print_progress('Writing output into %s.' % args.outdir)
os.makedirs(args.outdir, exist_ok=True)
train_dataset = None
if not args.gen:
# Load dataset for classifier case.
# Load perturbation
perturbation = numpy.load(
args.perturbation) if args.perturbation else None
segrunner = None
# Load broden dataset
if args.imagedir is not None:
dataset = try_to_load_images(args.imagedir, args.imgsize,
perturbation, args.size)
segrunner = ImageOnlySegRunner(dataset)
else:
dataset = try_to_load_broden(args.segments, args.imgsize, 1,
perturbation, args.download,
args.size)
if dataset is None:
dataset = try_to_load_multiseg(args.segments, args.imgsize,
perturbation, args.size)
if dataset is None:
print_progress('No segmentation dataset found in %s',
args.segments)
print_progress('use --download to download Broden.')
sys.exit(1)
else:
# For segmenter case the dataset is just a random z
#dataset = z_dataset_for_model(model, args.size)
#train_dataset = z_dataset_for_model(model, args.size, seed=2)
dataset = datasets.ImageFolder('dataset/Adissect',
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
(0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
]))
train_dataset = dataset
segrunner = GeneratorSegRunner(autoimport_eval(args.segmenter))
torch.no_grad()
# Run dissect
dissect(args.outdir,
model,
dataset,
train_dataset=train_dataset,
segrunner=segrunner,
examples_per_unit=args.examples,
netname=args.netname,
quantile_threshold=args.quantile_threshold,
meta=meta,
merge=mergedata,
make_images=args.images,
make_labels=args.labels,
make_maxiou=args.maxiou,
make_covariance=args.covariance,
make_report=args.report,
make_row_images=args.images,
make_single_images=True,
rank_all_labels=args.rank_all_labels,
batch_size=args.batch_size,
num_workers=args.num_workers,
settings=vars(args))
# Mark the directory so that it's not done again.
mark_job_done(args.outdir)