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)