def _switch(args, data, index, wolf, clabel):
print('switch: {}, #{}'.format(clabel, len(index)))
wolf.eval()
batch = 64
np.random.shuffle(index)
for run in range(5):
img0, y0 = get_batch(data, index[:batch])
img0 = img0.to(args.device)
y0 = y0.to(args.device)
img1, y1 = get_batch(data, index[(run + 1) * batch:(run + 2) * batch])
img1 = img1.to(args.device)
y1 = y1.to(args.device)
image_size = (3, args.image_size, args.image_size)
z0, epsilon0 = wolf.encode(img0,
y=y0,
n_bits=args.n_bits,
random=False)
z1, epsilon1 = wolf.encode(img1,
y=y1,
n_bits=args.n_bits,
random=False)
alphas = torch.arange(2, device=args.device).float().view(1, 2, 1)
# [1, time, 1, 1, 1]
betas = [0, 1]
# [batch, time, dim]
epsilon0 = epsilon0.expand(-1, alphas.size(1), *image_size)
epsilon1 = epsilon1.expand(-1, alphas.size(1), *image_size)
imgs = []
for beta in betas:
# [batch, time, c, h, w]
epsilon = epsilon0 * (1.0 - beta) + epsilon1 * beta
# [batch, time, dim]
z = z0 * (1.0 - alphas
) + z1 * alphas if beta == 0 else z0 * alphas + z1 * (
1.0 - alphas)
# [batch * time, *]
z = z.view(-1, z.size(2))
epsilon = epsilon.view(-1, *image_size)
# [batch, time, c, h, w]
img = wolf.decode(epsilon, z=z,
n_bits=args.n_bits).view(batch, -1, *image_size)
imgs.append(img)
nn = int(math.sqrt(batch))
# [batch, 2, 2, c, h, w]
img = torch.stack(imgs, dim=1)
# [nn, nn, 2, 2, c, h, w] -> [nn, 2, nn, 2, c, h, w]
img = img.view(nn, nn, 2, 2, *image_size).transpose(1, 2)
img = img.contiguous().view(-1, *image_size).cpu()
image_file = 'switch{}.png'.format(clabel + '-' + str(run))
save_image(img,
os.path.join(args.result_path, image_file),
nrow=2 * nn)
def reconstruct(args, epoch, val_data, val_index, wolf):
logging('reconstruct', args.log)
wolf.eval()
n = 16
np.random.shuffle(val_index)
img, y = get_batch(val_data, val_index[:n])
img = img.to(args.device)
y = y.to(args.device)
z, epsilon = wolf.encode(img, y=y, n_bits=args.n_bits, random=False)
epsilon = epsilon.squeeze(1)
z = z.squeeze(1) if z is not None else z
img_recon = wolf.decode(epsilon, z=z, n_bits=args.n_bits)
img = postprocess(preprocess(img, args.n_bits), args.n_bits)
abs_err = img_recon.add(img * -1).abs()
logging(
'Err: {:.4f}, {:.4f}'.format(abs_err.max().item(),
abs_err.mean().item()), args.log)
comparison = torch.cat([img, img_recon], dim=0).cpu()
reorder_index = torch.from_numpy(
np.array([[i + j * n for j in range(2)] for i in range(n)])).view(-1)
comparison = comparison[reorder_index]
image_file = 'reconstruct{}.png'.format(epoch)
save_image(comparison, os.path.join(args.result_path, image_file), nrow=16)
def reconstruct(args, data, wolf):
print('reconstruct')
wolf.eval()
batch = 16
nsamples = 15
index = np.arange(len(data))
np.random.shuffle(index)
img, y = get_batch(data, index[:batch])
img = img.to(args.device)
y = y.to(args.device)
image_size = (3, args.image_size, args.image_size)
_, epsilon = wolf.encode(img, y=y, n_bits=args.n_bits, nsamples=nsamples, random=True)
epsilon = epsilon.view(batch * nsamples, *image_size)
z = wolf.encode_global(img, y=y, n_bits=args.n_bits, nsamples=nsamples, random=True)
z = z.view(batch * nsamples, z.size(2))
# [batch, nsamples, c, h, w]
img_recon = wolf.decode(epsilon, z=z, n_bits=args.n_bits).view(batch, nsamples, *image_size)
# [batch, 1, c, h, w]
img = postprocess(preprocess(img, args.n_bits), args.n_bits).unsqueeze(1)
# [batch, nsamples + 1, c, h, w] -> [batch*(nsamples + 1), c, h, w]
comparison = torch.cat([img, img_recon], dim=1).view(-1, *image_size).cpu()
image_file = 'reconstruct.png'
save_image(comparison, os.path.join(args.result_path, image_file), nrow=nsamples + 1)
def init_model(args, train_data, train_index, wolf):
wolf.eval()
init_batch_size = args.init_batch_size
logging(
'Rank {}, init model: {} instances'.format(args.rank, init_batch_size),
args.log)
init_index = np.random.choice(train_index, init_batch_size, replace=False)
init_x, init_y = get_batch(train_data, init_index)
init_x = preprocess(init_x.to(args.device), args.n_bits)
init_y = init_y.to(args.device)
wolf.init(init_x, y=init_y, init_scale=1.0)
def switch(args, data, wolf, num_class):
index = np.arange(len(data))
_switch(args, data, index, wolf, 'g')
if num_class is not None:
index = np.arange(len(data))
_, y = get_batch(data, index)
index = torch.from_numpy(index)
for label in range(num_class):
mask = y.eq(label)
idx = index[mask].numpy()
_switch(args, data, idx, wolf, label)
def _interpolate(args, data, index, wolf, clabel):
print('interpolate: {}, #{}'.format(clabel, len(index)))
wolf.eval()
batch = 64
np.random.shuffle(index)
img0, y0 = get_batch(data, index[:batch])
img0 = img0.to(args.device)
y0 = y0.to(args.device)
img1, y1 = get_batch(data, index[batch:2 * batch])
img1 = img1.to(args.device)
y1 = y1.to(args.device)
image_size = (3, args.image_size, args.image_size)
z0, epsilon0 = wolf.encode(img0, y=y0, n_bits=args.n_bits, random=False)
z1, epsilon1 = wolf.encode(img1, y=y1, n_bits=args.n_bits, random=False)
alphas = [x * 0.1 for x in range(11)]
# [1, time, 1, 1, 1]
betas = torch.arange(11, device=args.device).float().view(1, 11, 1, 1,
1) * 0.1
# [batch, time, dim]
z0 = z0.expand(-1, betas.size(1), -1)
z1 = z1.expand(-1, betas.size(1), -1)
imgs = []
for alpha in alphas:
# [batch, time, dim]
z = z0 * (1.0 - alpha) + z1 * alpha
# [batch, time, c, h, w]
epsilon = epsilon0 * (1.0 - betas) + epsilon1 * betas
# [batch * time, *]
z = z.view(-1, z.size(2))
epsilon = epsilon.view(-1, *image_size)
# [batch, time, c, h, w]
img = wolf.decode(epsilon, z=z,
n_bits=args.n_bits).view(batch, -1, *image_size)
imgs.append(img)
img = torch.stack(imgs, dim=1).view(-1, *image_size).cpu()
image_file = 'interpolate{}.png'.format(clabel)
save_image(img, os.path.join(args.result_path, image_file), nrow=11)