def sample_latent(self, n_samples=1, seed=None, truncation=None):
if seed is None:
seed = np.random.randint(np.iinfo(
np.int32).max) # use (reproducible) global rand state
noise = zdataset.z_sample_for_model(self.model, n_samples,
seed=seed)[...]
return noise.to(self.device)
def main():
expdir = 'results/pgw/%s/%s/layer%d' % (ganname, modelname, layernum)
if ganname == 'proggan':
model = setting.load_proggan(modelname).cuda()
zds = zdataset.z_dataset_for_model(model, size=1000)
writer = ganrewrite.ProgressiveGanRewriter
elif ganname == 'stylegan':
model = load_seq_stylegan(modelname, mconv='seq')
zds = zdataset.z_dataset_for_model(model, size=1000)
writer = ganrewrite.SeqStyleGanRewriter
model.eval()
gw = writer(model, zds, layernum, cachedir=expdir)
images = []
with torch.no_grad():
for _ in tqdm(range(N//batch_size + 1)):
z = zdataset.z_sample_for_model(model, size=batch_size, seed=len(images)).cuda()
samples = gw.sample_image_patch(z, crop_size)
samples = [s.data.cpu() for s in samples]
images.extend(samples)
images = torch.stack(images[:N], dim=0)
gt_fid = 0
fake_fid = compute_fid(images, f'{modelname}_cropped_{images.size(2)}_{ganname}')
save_image(images[:32] * 0.5 + 0.5, f'patches_{layernum}_{ganname}_{modelname}_{crop_size}.png')
return fake_fid, gt_fid, images.size(2)
def standard_z_sample(self, size=100, seed=1, device=None):
'''
Generate a standard set of random Z as a (size, z_dimension) tensor.
With the same random seed, it always returns the same z (e.g.,
the first one is always the same regardless of the size.)
'''
result = z_sample_for_model(self.model, size)
if device is not None:
result = result.to(device)
return result
def seeded_cropped_sample(g,
gw,
imgnum,
crop_seed,
crop_size,
act=True,
size=None):
with torch.no_grad():
z = zdataset.z_sample_for_model(g, size=1, seed=imgnum).cuda()
return gw.sample_image_patch(z,
crop_size,
seed=crop_seed,
act=act,
size=size)
def get_samples(model, name, nimgs=50000, truncated=False):
batch_size = 10
g = load_model(model, name, truncated)
g.eval()
with torch.no_grad():
samples = []
for _ in tqdm(range(nimgs // batch_size + 1)):
seed = len(samples) if samples is not None else 0
z = zdataset.z_sample_for_model(g, size=batch_size,
seed=seed).cuda()
x_real = g(z)
x_real = [x.detach().cpu() for x in x_real]
samples.extend(x_real)
samples = torch.stack(samples, dim=0)
return samples
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)