def sample_images(batches_done):
"""Saves a generated sample from the test set"""
imgs = next(iter(val_dataloader))
G_AB.eval()
G_BA.eval()
Trans_AB.eval()
Trans_BA.eval()
real_A_full = Variable(imgs["A"].type(Tensor))
pyr_A = pyramid.pyramid_decom(img=real_A_full, max_levels=2)
real_A = pyr_A[-1]
fake_B = G_AB(real_A)
pyr_A_trans = pyramid_transform(pyr_A, real_A, fake_B, Trans_AB)
fake_B_full = pyramid.pyramid_recons(pyr_A_trans)
real_B_full = Variable(imgs["B"].type(Tensor))
pyr_B = pyramid.pyramid_decom(img=real_B_full, max_levels=2)
real_B = pyr_B[-1]
fake_A = G_BA(real_B)
pyr_B_trans = pyramid_transform(pyr_B, real_B, fake_A, Trans_BA)
fake_A_full = pyramid.pyramid_recons(pyr_B_trans)
# Arange images along x-axis
real_A = make_grid(real_A_full.cpu().data, nrow=5, normalize=True)
real_B = make_grid(real_B_full.cpu().data, nrow=5, normalize=True)
fake_A = make_grid(fake_A_full.cpu().data, nrow=5, normalize=True)
fake_B = make_grid(fake_B_full.cpu().data, nrow=5, normalize=True)
# Arange images along y-axis
image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1)
save_image(image_grid, "images/%s/%s.png" % (opt.dataset_name, batches_done), normalize=False)
def sample_images(batches_done):
"""Saves a generated sample from the test set"""
imgs = next(iter(val_dataloader))
G.eval()
Trans.eval()
real_full = Variable(imgs["A"].type(Tensor))
real_B_full = Variable(imgs["B"].type(Tensor))
start = time.time()
pyr = pyramid.pyramid_decom(img=real_full, max_levels=opt.levels)
fake = G(pyr[-1])
pyr_trans = pyramid_transform(pyr, pyr[-1], fake, Trans, if_conv)
fake_full = pyramid.pyramid_recons(pyr_trans)
cost = (time.time() - start) / opt.validate_size
print('time cost for one image: {:.4f}'.format(cost))
# Arange images along x-axis
real_A = make_grid(torch.clamp(real_full, -1, 1).cpu().data, nrow=5, normalize=True)
fake_B = make_grid(torch.clamp(fake_full, -1, 1).cpu().data, nrow=5, normalize=True)
real_B = make_grid(torch.clamp(real_B_full, -1, 1).cpu().data, nrow=5, normalize=True)
# Arange images along y-axis
image_grid = torch.cat((real_A, fake_B, real_B), 1)
if not os.path.exists("../images/%s" % (opt.dataset_name)):
os.makedirs("../images/%s" % (opt.dataset_name))
# mse = torch.nn.functional.mse_loss(torch.clamp(fake_full[0], -1, 1), torch.clamp(real_B_full[0], -1, 1))
# print(fake_full[0].shape, real_B_full[0].shape)
# psnr = 10 * log10(1 / mse.item())
# print('psnr: {:.4f}'.format(psnr))
save_image(image_grid, "../images/%s/%s.png" % (opt.dataset_name, batches_done), normalize=False)
def sample_images(batches_done):
"""Saves a generated sample from the test set"""
imgs = next(iter(val_dataloader))
G_AB.eval()
G_BA.eval()
Trans_AB.eval()
Trans_BA.eval()
real_A_full = Variable(imgs["A"].type(Tensor))
pyr_A = pyramid.pyramid_decom(img=real_A_full, max_levels=opt.levels)
real_A = pyr_A[-1]
fake_B = G_AB(real_A)
pyr_A_trans = pyramid_transform(pyr_A, real_A, fake_B, Trans_AB, if_conv)
fake_B_full = pyramid.pyramid_recons(pyr_A_trans)
real_B_full = Variable(imgs["B"].type(Tensor))
start = time.time()
pyr_B = pyramid.pyramid_decom(img=real_B_full, max_levels=opt.levels)
real_B = pyr_B[-1]
fake_A = G_BA(real_B)
pyr_B_trans = pyramid_transform(pyr_B, real_B, fake_A, Trans_BA, if_conv)
fake_A_full = pyramid.pyramid_recons(pyr_B_trans)
cost = (time.time() - start) / opt.validate_size
print('time cost for one image: {:.4f}'.format(cost))
# Arange images along x-axis
real_A = make_grid(torch.clamp(real_A_full, -1, 1).cpu().data,
nrow=5,
normalize=True)
real_B = make_grid(torch.clamp(real_B_full, -1, 1).cpu().data,
nrow=5,
normalize=True)
fake_A = make_grid(torch.clamp(fake_A_full, -1, 1).cpu().data,
nrow=5,
normalize=True)
fake_B = make_grid(torch.clamp(fake_B_full, -1, 1).cpu().data,
nrow=5,
normalize=True)
# Arange images along y-axis
image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1)
save_image(image_grid,
"../images/%s/%s.png" % (opt.dataset_name, batches_done),
normalize=False)
normalize=False)
# ----------
# Training
# ----------
for epoch in range(opt.epoch, opt.n_epochs):
for i, batch in enumerate(dataloader):
start_time = time.time()
# Set model input
real_A_full = Variable(batch["A"].type(Tensor))
real_B_full = Variable(batch["B"].type(Tensor))
pyr_A = pyramid.pyramid_decom(img=real_A_full, max_levels=opt.levels)
pyr_B = pyramid.pyramid_decom(img=real_B_full, max_levels=opt.levels)
# save_image(real_A, "1.png", normalize=False)
# save_image(image_recons, "2.png", normalize=False)
# print(image_recons.shape)
# print(real_A.shape)
real_A = pyr_A[-1]
real_B = pyr_B[-1]
# Adversarial ground truths
valid = Variable(Tensor(np.ones((real_A.size(0), *D_A.output_shape))),
requires_grad=False)
fake = Variable(Tensor(np.zeros((real_A.size(0), *D_A.output_shape))),
requires_grad=False)
