test_loader = util.data_load('data/' + opt.dataset,
opt.test_subfolder,
transform,
opt.test_batch_size,
shuffle=True)
test = test_loader.__iter__().__next__()[0]
img_size = test.size()[2]
if opt.inverse_order:
fixed_y_ = test[:, :, :, 0:img_size]
fixed_x_ = test[:, :, :, img_size:]
else:
fixed_x_ = test[:, :, :, 0:img_size]
fixed_y_ = test[:, :, :, img_size:]
if img_size != opt.input_size:
fixed_x_ = util.imgs_resize(fixed_x_, opt.input_size)
fixed_y_ = util.imgs_resize(fixed_y_, opt.input_size)
# network
G = network.generator(opt.ngf)
D = network.discriminator(opt.ndf)
G.weight_init(mean=0.0, std=0.02)
D.weight_init(mean=0.0, std=0.02)
G.train()
D.train()
# loss
BCE_loss = nn.BCELoss()
L1_loss = nn.L1Loss()
# Adam optimizer
G.cuda()
G.load_state_dict(torch.load(opt.dataset + '_results/' + opt.dataset + '_generator_param.pkl'))
# network
n = 0
print('test start!')
for x_, _ in test_loader:
if opt.inverse_order:
y_ = x_[:, :, :, :x_.size()[2]]
x_ = x_[:, :, :, x_.size()[2]:]
else:
y_ = x_[:, :, :, x_.size()[2]:]
x_ = x_[:, :, :, :x_.size()[2]]
if x_.size()[2] != opt.input_size:
x_ = util.imgs_resize(x_, opt.input_size)
y_ = util.imgs_resize(y_, opt.input_size)
x_ = Variable(x_.cuda(), volatile=True)
test_image = G(x_)
s = test_loader.dataset.imgs[n][0][::-1]
s_ind = len(s) - s.find('/')
e_ind = len(s) - s.find('.')
ind = test_loader.dataset.imgs[n][0][s_ind:e_ind-1]
path = opt.dataset + '_results/test_results/' + ind + '_input.png'
plt.imsave(path, (x_[0].cpu().data.numpy().transpose(1, 2, 0) + 1) / 2)
path = opt.dataset + '_results/test_results/' + ind + '_output.png'
plt.imsave(path, (test_image[0].cpu().data.numpy().transpose(1, 2, 0) + 1) / 2)
path = opt.dataset + '_results/test_results/' + ind + '_target.png'
plt.imsave(path, (y_[0].numpy().transpose(1, 2, 0) + 1) / 2)
D_A_optimizer.param_groups[0]['lr'] -= opt.lrD / (opt.train_epoch -
opt.decay_epoch)
D_B_optimizer.param_groups[0]['lr'] -= opt.lrD / (opt.train_epoch -
opt.decay_epoch)
G_optimizer.param_groups[0]['lr'] -= opt.lrG / (opt.train_epoch -
opt.decay_epoch)
#H_A_optimizer.param_groups[0]['lr'] = opt.lrH*(0.1**(epoch//opt.decay_epoch))
start_batch_time = time.time()
for iteration, batch in enumerate(train_loader, 0):
realA = batch[0]
realB = batch[1]
batch_ind = batch[2]
if opt.resize_scale:
realA = util.imgs_resize(realA, opt.resize_scale)
realB = util.imgs_resize(realB, opt.resize_scale)
if opt.crop:
realA = util.random_crop(realA, opt.input_size)
realB = util.random_crop(realB, opt.input_size)
if opt.fliplr:
realA = util.random_fliplr(realA)
realB = util.random_fliplr(realB)
realA, realB = Variable(realA.cuda()), Variable(realB.cuda())
# G STEP
# train generator G