def evaluate(self, task, model, seed):
model = _check_model(model)
model.eval()
stats = OrderedDict()
itr = task.get_batch_iterator(dataset=task.dataset,
batch_size=self.args.batch_size,
num_batches=self.args.eval_batches,
seed=seed).next_epoch_itr(shuffle=False)
zs = []
ys = []
for batch in itr:
if self.cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
# (batch_size, code_size)
zs.append(outputs['z'].detach().cpu().numpy())
# (batch_size, num_factors)
ys.append(batch['factor'].detach().cpu().numpy())
# (code_size, eval_batches x batch_size)
zs = np.transpose(np.concatenate(zs, axis=0))
# (num_factors, eval_batches x batch_size)
ys = np.transpose(np.concatenate(ys, axis=0))
discrete_zs = _discretize(zs, self.args.num_bins)
# (code_size, num_factors)
mi = _discrete_mutual_info(discrete_zs, ys)
assert mi.shape[0] == zs.shape[0]
assert mi.shape[1] == ys.shape[0]
stats['modularity'] = modularity(mi)
return stats
def _gather_votes(task, model, num_batches, batch_size, global_variances,
active_dims, seed, cuda):
num_factors = task.dataset.num_factors
all_factors = np.arange(num_factors)
code_size = global_variances.shape[0]
votes = np.zeros((num_factors, code_size), dtype=int)
for factor_index in range(num_factors):
fixed_factors = np.delete(all_factors, factor_index)
itr = task.get_iterator_with_fixed_factor(
dataset=task.dataset,
batch_size=batch_size,
num_batches=num_batches,
factor_index=factor_index,
seed=seed + 1 + factor_index).next_epoch_itr(shuffle=False)
for batch in itr:
if cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
reprs = outputs['z'].detach().cpu().numpy()
local_variances = np.var(reprs, axis=0, ddof=1)
code_index = np.argmax(local_variances[active_dims] /
global_variances[active_dims])
votes[factor_index, code_index] += 1
return votes
def evaluate(self, task, model, seed):
model = _check_model(model)
model.eval()
stats = OrderedDict()
itr = task.get_batch_iterator(dataset=task.dataset,
batch_size=self.args.batch_size,
num_batches=self.args.eval_batches,
seed=seed).next_epoch_itr(shuffle=False)
zs = []
for batch in itr:
if self.cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
# (batch_size, code_size)
zs.append(outputs['z'].detach().cpu().numpy())
# (code_size, eval_batches x batch_size)
zs = np.transpose(np.concatenate(zs, axis=0))
cov = np.cov(zs)
stats['total_correlation'] = gaussian_total_correlation(cov)
stats['wasserstein_correlation'] = gaussian_wasserstein_correlation(
cov)
stats['wasserstein_correlation_norm'] = stats['wasserstein_correlation'] \
/ np.sum(np.diag(cov))
return stats
def evaluate(self, task, model, seed):
model = _check_model(model)
model.eval()
stats = OrderedDict()
itr = task.get_batch_iterator(dataset=task.dataset,
batch_size=self.args.batch_size,
num_batches=self.args.eval_batches,
seed=seed).next_epoch_itr(shuffle=False)
zs = []
ys = []
for batch in itr:
if self.cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
# (batch_size, code_size)
zs.append(outputs['z'].detach().cpu().numpy())
# (batch_size, num_factors)
ys.append(batch['factor'].detach().cpu().numpy())
# (code_size, eval_batches x batch_size)
zs = np.transpose(np.concatenate(zs, axis=0))
# normalize z to have scale 1
zs /= zs.max(1, keepdims=True) - zs.min(1, keepdims=True)
np.random.seed(seed)
shuffled_zs = [np.random.permutation(z) for z in zs]
# (num_factors, eval_batches x batch_size)
ys = np.transpose(np.concatenate(ys, axis=0))
wmis = []
for factor_idx, yi in enumerate(ys):
wmi = []
# calculate W1 for each value of the given factor
for y in np.unique(yi):
mask = yi == y
# code_size
wmi.append(
[wasserstein_distance(
z[mask], shuffled_z[mask]) \
for z, shuffled_z in zip(zs, shuffled_zs)]
)
# this becomes (num_factors, code_size)
wmis.append(np.mean(wmi, axis=0))
# (code_size, num_factors)
wmis = np.transpose(wmis)
sorted_wmi = np.sort(wmis, axis=0)[::-1]
wmig = sorted_wmi[0, :] - sorted_wmi[1, :]
for i in range(len(wmig)):
stats['wmig_factor_{}'.format(i)] = wmig[i]
avg_wmig = np.mean(wmig[np.isfinite(wmig) & np.greater(wmig, 0)])
stats['avg_wmig'] = avg_wmig
return stats
def _estimate_variances(task, model, num_batches, batch_size, seed, cuda):
epoch_iter = task.get_batch_iterator(
dataset=task.dataset,
batch_size=batch_size,
num_batches=num_batches,
seed=seed).next_epoch_itr(shuffle=False)
code_samples = []
for batch in epoch_iter:
if cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
# z has size (batch_size, code_size)
code_samples.append(outputs['z'].detach().cpu().numpy())
code_samples = np.concatenate(code_samples, axis=0)
return np.var(code_samples, axis=0, ddof=1)
def evaluate(self, task, model, seed):
model = _check_model(model)
model.eval()
stats = OrderedDict()
itr = task.get_batch_iterator(dataset=task.dataset,
batch_size=self.args.batch_size,
num_batches=self.args.eval_batches,
seed=seed).next_epoch_itr(shuffle=False)
rec_meter = AverageMeter()
for batch in itr:
if self.cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
batch_size = batch['batch_size']
rec = torch.sum(
-outputs['x'].log_prob(batch['image'])) / batch_size
rec_meter.update(rec.item(), batch_size)
stats['reconstruction'] = rec_meter.avg
return stats
def evaluate(self, task, model, seed):
model = _check_model(model)
model.eval()
stats = OrderedDict()
itr = task.get_batch_iterator(dataset=task.dataset,
batch_size=self.args.batch_size,
num_batches=self.args.eval_batches,
seed=seed).next_epoch_itr(shuffle=False)
zs = []
ys = []
for batch in itr:
if self.cuda:
batch = utils.move_to_cuda(batch)
outputs = model(batch)
# (batch_size, code_size)
zs.append(outputs['z'].detach().cpu().numpy())
# (batch_size, num_factors)
ys.append(batch['factor'].detach().cpu().numpy())
# (code_size, eval_batches x batch_size)
zs = np.transpose(np.concatenate(zs, axis=0))
# (num_factors, eval_batches x batch_size)
ys = np.transpose(np.concatenate(ys, axis=0))
discrete_zs = _discretize(zs, self.args.num_bins)
# (code_size, num_factors)
mi = _discrete_mutual_info(discrete_zs, ys)
# (num_factors,)
entropy = _discrete_entropy(ys)
sorted_mi = np.sort(mi, axis=0)[::-1]
with np.errstate(divide='ignore', invalid='ignore'):
mig = np.divide(sorted_mi[0, :] - sorted_mi[1, :], entropy[:])
for i in range(len(mig)):
stats['mig_factor_{}'.format(i)] = mig[i]
stats['avg_mig'] = np.mean(mig[np.isfinite(mig)])
return stats
def _prepare_sample(self, sample):
if sample is None or len(sample) == 0:
return None
if self.cuda:
sample = utils.move_to_cuda(sample)
return sample
def main(args):
assert args.path is not None, '--path required for generation'
print(args)
use_cuda = torch.cuda.is_available() and not args.cpu
if use_cuda:
torch.cuda.set_device(args.device_id)
task = tasks.setup_task(args)
task.load_dataset()
print('| loading model from {}'.format(args.path))
model, model_args = utils.load_model_for_inference(args.path, task)
if use_cuda:
model.cuda()
itr = task.get_batch_iterator(dataset=task.dataset,
batch_size=args.batch_size,
num_batches=args.gen_batches,
seed=args.seed +
1).next_epoch_itr(shuffle=False)
gen_timer = StopwatchMeter()
generator = task.build_generator(args)
pb = progress_bar.build_progress_bar(args, itr)
mod_values = get_modification_values(args)
for batch in pb:
batch = utils.move_to_cuda(batch) if use_cuda else batch
batch_size = batch['batch_size']
if args.gen_mode == 'standard' or args.gen_mode is None:
dims = []
elif args.modify_dims is None:
dims = range(model_args.code_size)
else:
dims = args.modify_dims
mods = [{'dim': idx, 'values': mod_values} for idx in dims]
gen_timer.start()
batch_images = task.gen_step(generator, model, batch, mods)
if len(mods) == 0:
# (1, B)
num_images = batch_size
args.save_format = 'jpeg'
save_batch_images(batch['id'], batch_images[0], args)
else:
# (D, V, B)
num_images = batch_size * len(mod_values) * len(mods)
for i, dim in enumerate(dims):
if args.save_format.lower() == 'gif':
for j in range(batch_size):
save_gifs(batch['id'][j],
batch_images[i, :, j],
args,
dim=dim)
else:
for j, val in enumerate(mod_values):
save_batch_images(batch['id'],
batch_images[i][j],
args,
dim=dim,
mod_val=val)
gen_timer.stop(num_images)
print('| generated {} images in {:.1f}s ({:.2f} images/s)'.format(
gen_timer.n, gen_timer.sum, gen_timer.n / gen_timer.sum))
