def __init__(self, pthdir, layers=None, device=None, dissectdir=None):
os.makedirs(pthdir, exist_ok=True)
self.device = device if device is not None else torch.device('cpu')
self.dissectdir = dissectdir if dissectdir is not None else (
os.path.join(pthdir, 'dissect'))
self.modellock = threading.Lock()
# Load the generator from the pth file. If the file is not there, download
path_gan_checkpoint = os.path.join(pthdir, 'generator.pth')
if not os.path.isfile(path_gan_checkpoint):
wget.download(
'http://wednesday.csail.mit.edu/gaze/ganclevr/files/generator.pth',
out=path_gan_checkpoint)
model = proggan.from_pth_file(
path_gan_checkpoint,
map_location=lambda storage, location: storage)
model.eval()
self.model = model
# Get the set of layers of interest.
# Default: all shallow children except last.
if layers is None:
layers = [name for name, module in model.named_children()][:-1]
self.layers = layers
# Modify model to instrument the given layers
retain_layers(model, layers)
edit_layers(model, layers)
# Move it to CUDA if wanted.
model.to(device)
# Determine z dimension.
self.z_dimension = [
c for c in model.modules() if isinstance(c, torch.nn.Conv2d)
][0].in_channels
# Run the model on one sample input to determine output image size as well as feature size of every layer
z = torch.randn(self.z_dimension)[None, :, None, None].to(device)
output = model(z)
self.image_shape = output.shape[2:]
self.layer_shape = {
layer: tuple(model.retained[layer].shape)
for layer in layers
}
for param in self.model.parameters():
param.requires_grad = False
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 main():
parser = argparse.ArgumentParser(description='GAN sample making utility')
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='images',
help='directory for image output')
parser.add_argument('--size', type=int, default=100,
help='number of images to output')
parser.add_argument('--test_size', type=int, default=None,
help='number of images to test')
parser.add_argument('--layer', type=str, default=None,
help='layer to inspect')
parser.add_argument('--seed', type=int, default=1,
help='seed')
parser.add_argument('--maximize_units', type=int, nargs='+', default=None,
help='units to maximize')
parser.add_argument('--ablate_units', type=int, nargs='+', default=None,
help='units to ablate')
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()
verbose_progress(not args.quiet)
# Instantiate the model
model = autoimport_eval(args.model)
if args.pthfile 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)
# Unwrap any DataParallel-wrapped model
if isinstance(model, torch.nn.DataParallel):
model = next(model.children())
# 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)):
z_channels = first_layer.in_channels
spatialdims = (1, 1)
else:
z_channels = first_layer.in_features
spatialdims = ()
# Instrument the model if needed
if args.maximize_units is not None:
retain_layers(model, [args.layer])
model.cuda()
# Get the sample of z vectors
if args.maximize_units is None:
indexes = torch.arange(args.size)
z_sample = standard_z_sample(args.size, z_channels, seed=args.seed)
z_sample = z_sample.view(tuple(z_sample.shape) + spatialdims)
else:
# By default, if maximizing units, get a 'top 5%' sample.
if args.test_size is None:
args.test_size = args.size * 20
z_universe = standard_z_sample(args.test_size, z_channels,
seed=args.seed)
z_universe = z_universe.view(tuple(z_universe.shape) + spatialdims)
indexes = get_highest_znums(model, z_universe, args.maximize_units,
args.size, seed=args.seed)
z_sample = z_universe[indexes]
if args.ablate_units:
edit_layers(model, [args.layer])
dims = max(2, max(args.ablate_units) + 1) # >=2 to avoid broadcast
model.ablation[args.layer] = torch.zeros(dims)
model.ablation[args.layer][args.ablate_units] = 1
save_znum_images(args.outdir, model, z_sample, indexes,
args.layer, args.ablate_units)
copy_lightbox_to(args.outdir)