def hybridize(self, q): hyb = util.shuffle_dim(q) if self.shuffle: return hyb sel = (torch.rand_like(q) - self.probs).gt(0).float() return q*sel + hyb*(1-sel)
def generate(N):
idx = torch.randperm(len(Q))[:N]
q = Q[idx].to(A.device)
with torch.no_grad():
gen = model.decode(util.shuffle_dim(q)).detach()
return gen
def gen_target(model, X=None, Q=None, hybrid=False, ret_q=False):
if Q is None:
assert X is not None
Q = model.encode(X)
if isinstance(Q, distrib.Distribution):
Q = Q.mean
if hybrid:
Q = util.shuffle_dim(Q)
gen = model.decode(Q)
if ret_q:
return gen, Q
return gen
def get_stats(model, model_inception, batch_size, n_samples, dataset, mode = True, hybrid = False, stat_ref = False, save_stats = None):
pred_arr = np.empty((n_samples, 2048))
out = {}
if stat_ref:
if dataset == '3dshapes':
data = datasets.Shapes3D(dataroot = os.environ['FOUNDATION_DATA_DIR'], train = mode, labels = True)
else:
raise NotImplementedError
assert len(data) >= n_samples, 'Number of samples should be less than dataset size'
indices = np.random.randint(0, len(data), n_samples)
j = 0
print('Computing Inception Features')
while j < n_samples:
if j% (batch_size*100) == 0:
print('Done', j , 'Samples')
curr_batch_size = min(batch_size, n_samples - j)
with torch.no_grad():
if stat_ref:
generated, _ = data[indices[j:j+curr_batch_size]]
generated = generated.to(model.device)
else:
if hybrid:
q = model.sample_prior(curr_batch_size)
generated = model.decode(util.shuffle_dim(q))
else:
generated = model.generate(curr_batch_size)
pred = model_inception(generated)[0]
if pred.shape[2] != 1 or pred.shape[3] != 1:
pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
pred_arr[j:j+curr_batch_size] = pred.cpu().data.numpy().reshape(curr_batch_size, -1)
j += curr_batch_size
m = np.mean(pred_arr, axis=0)
s = np.cov(pred_arr, rowvar=False)
if save_stats is not None:
out['indices'] = indices
out['m'] = m
out['sigma'] = s
with open(save_stats,'wb') as b:
pkl.dump(out, b)
return m, s
def regularize(self, q):
qdis = q
if isinstance(q, distrib.Distribution):
q = q.loc
mix = util.shuffle_dim(q)
if self.latent_disc is None:
reg = util.MMD(q, mix)
else:
self.reg_step_counter += 1
vreal = self.latent_disc(mix)
vfake = self.latent_disc(q)
wasserstein = vreal.mean() - vfake.mean()
self.stats.update('factor-ws', wasserstein.detach())
loss = -wasserstein
if self.latent_disc_gp is not None and self.latent_disc_gp > 0:
lgp_loss = unsup.grad_penalty(self.latent_disc, mix, q)
self.stats.update('factor-gp', lgp_loss.detach())
loss += self.latent_disc_gp*lgp_loss
if self.train_me():
self.optim.latent_disc.zero_grad()
loss.backward(retain_graph=True)
self.optim.latent_disc.step()
if self.latent_disc_steps <= 0 or self.reg_step_counter % self.latent_disc_steps == 0:
reg = self.latent_disc(q).mean()
else:
reg = 0.
if self.prior_wt is not None and self.prior_wt > 0:
reg_prior = super().regularize(qdis)
self.stats.update('reg-prior', reg_prior)
self.stats.update('reg-factor', reg)
reg = (1-self.prior_wt)*reg + self.prior_wt*reg_prior
return reg
def run_model(S, pbar=None, **unused):
A = S.A
dataset = S.dataset
model = S.model
assert False, 'unused'
if 'loader' not in S:
# DataLoader
A.dataset.batch_size = 16
util.set_seed(0)
loader = train.get_loaders(
dataset,
batch_size=A.dataset.batch_size,
num_workers=A.num_workers,
shuffle=True,
drop_last=False,
)
util.set_seed(0)
loader = iter(loader)
S.loader = loader
common_Ws = {64: 8, 32: 4, 16: 4, 9: 3, 8: 2, 4: 2}
border, between = 0.02, 0.01
img_W = common_Ws[A.dataset.batch_size]
S.img_W = img_W
S.border, S.between = border, between
if 'X' not in S:
# Batch
batch = next(loader)
batch = util.to(batch, A.device)
S.batch = batch
S.X = batch[0]
X = S.X
with torch.no_grad():
if model.enc is None:
q = model.sample_prior(X.size(0))
else:
q = model.encode(X)
qdis = None
qmle = q
if isinstance(q, distrib.Distribution):
qdis = q
q = q.rsample()
qmle = qdis.loc
rec = model.decode(q)
# vrec = model.disc(rec) if model.disc is not None else None
p = model.sample_prior(X.size(0))
gen = model.decode(p)
h = util.shuffle_dim(q)
hyb = model.decode(h)
S.rec = rec
S.gen = gen
S.hyb = hyb
S.q = q
S.p = p
S.h = h
S.qdis = qdis
S.qmle = qmle
batch_size = 128 # number of samples to get distribution
util.set_seed(0)
int_batch = next(
iter(
train.get_loaders(
dataset,
batch_size=batch_size,
num_workers=A.num_workers,
shuffle=True,
drop_last=False,
)))
with torch.no_grad():
int_batch = util.to(int_batch, A.device)
int_X, = int_batch
if model.enc is None:
int_q = model.sample_prior(int_X.size(0))
else:
int_q = model.encode(int_X)
dis_int_q = None
if isinstance(int_q, distrib.Distribution):
# int_q = int_q.rsample()
dis_int_q = int_q
int_q = int_q.loc
del int_batch
del int_X
S.int_q = int_q
S.dis_int_q = dis_int_q
latent_dim = A.model.latent_dim
iH, iW = X.shape[-2:]
rows = 4
steps = 60
if 'bounds' in S:
bounds = S.bounds
else:
# bounds = -2,2
bounds = None
# dlayout = rows, latent_dim // rows
dlayout = util.calc_tiling(latent_dim)
outs = []
all_diffs = []
inds = [0, 2, 3]
inds = np.arange(len(q))
save_inds = [0, 1, 2, 3]
# save_inds = []
saved_walks = []
for idx in inds:
walks = []
for dim in range(latent_dim):
dev = int_q[:, dim].std()
if bounds is None:
deltas = torch.linspace(int_q[:, dim].min(),
int_q[:, dim].max(), steps)
else:
deltas = torch.linspace(bounds[0], bounds[1], steps)
vecs = torch.stack([int_q[idx]] * steps)
vecs[:, dim] = deltas
with torch.no_grad():
walks.append(model.decode(vecs).cpu())
walks = torch.stack(walks, 1)
chn = walks.shape[2]
dsteps = 10
diffs = (walks[dsteps:] - walks[:-dsteps]
).abs() # .view(steps-dsteps, latent_dim, chn, 64*64)
# diffs /= (walks[dsteps:]).abs()
# diffs = diffs.clamp(min=1e-10,max=1).abs()
diffs = diffs.view(steps - dsteps, latent_dim, chn * iH * iW).mean(-1)
# diffs = 1 - diffs.mean(-1)
# print(diffs.shape)
# diffs *= 2
all_diffs.append(diffs.mean(0))
# print(all_diffs[-1])
if idx in save_inds:
# save_dir = S.save_dir
walks_full = walks.view(steps, dlayout[0], dlayout[1], chn, iH, iW) \
.permute(0, 1, 4, 2, 5, 3).contiguous().view(steps, iH * dlayout[0], iW * dlayout[1], chn).squeeze()
images = []
for img in walks_full.cpu().numpy():
images.append((img * 255).astype(np.uint8))
saved_walks.append(images)
# imageio.mimsave(os.path.join(save_dir, 'walks-idx{}.gif'.format(idx, dim)), images)
#
# with open(os.path.join(save_dir, 'walks-idx{}.gif'.format(idx, dim)), 'rb') as f:
# outs.append(display.Image(data=f.read(), format='gif'))
del walks_full
all_diffs = torch.stack(all_diffs)
S.all_diffs = all_diffs
S.saved_walks = saved_walks
dataset = S.dataset
full_q = None
if model.enc is not None:
N = min(1000, len(dataset))
print('Using {} samples'.format(N))
assert (N // 4) * 4 == N, 'invalid num: {}'.format(N)
loader = train.get_loaders(
dataset,
batch_size=N // 4,
num_workers=A.num_workers,
shuffle=True,
drop_last=False,
)
util.set_seed(0)
if pbar is not None:
loader = pbar(loader, total=len(loader))
loader.set_description('Collecting latent vectors')
full_q = []
for batch in loader:
batch = util.to(batch, A.device)
X, = batch
with torch.no_grad():
q = model.encode(X)
if isinstance(q, distrib.Distribution):
q = torch.stack([q.loc, q.scale], 1)
full_q.append(q.cpu())
if pbar is not None:
loader.close()
if len(full_q):
full_q = torch.cat(full_q)
# print(full_q.shape)
if len(full_q.shape) > 2:
full_q = distrib.Normal(loc=full_q[:, 0], scale=full_q[:, 1])
else:
full_q = None
S.full_q = full_q
if full_q is not None:
if 'results' not in S:
S.results = {}
S.results['val_Q'] = full_q
print('Storing {} latent vectors'.format(
len(full_q if not isinstance(full_q, distrib.Distribution) else
full_q.loc)))