y_ = x_[:, :, :, img_size:]
x_ = x_[:, :, :, 0:img_size]
if img_size != opt.input_size:
x_ = util.imgs_resize(x_, opt.input_size)
y_ = util.imgs_resize(y_, opt.input_size)
if opt.resize_scale:
x_ = util.imgs_resize(x_, opt.resize_scale)
y_ = util.imgs_resize(y_, opt.resize_scale)
if opt.crop_size:
x_, y_ = util.random_crop(x_, y_, opt.crop_size)
if opt.fliplr:
x_, y_ = util.random_fliplr(x_, y_)
x_, y_ = Variable(x_), Variable(y_)
D_result = D(x_, y_).squeeze()
D_real_loss = BCE_loss(D_result, Variable(torch.ones(D_result.size())))
G_result = G(x_)
D_result = D(x_, G_result).squeeze()
D_fake_loss = BCE_loss(D_result,
Variable(torch.zeros(D_result.size())))
D_train_loss = (D_real_loss + D_fake_loss) * 0.5
D_train_loss.backward()
D_optimizer.step()
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
G_optimizer.zero_grad()
# generate real A to fake B; D_A(G_A(A))
fakeB = G_A(realA)
D_A_result = D_A(fakeB)
G_A_loss = MSE_loss(D_A_result,
Variable(torch.ones(D_A_result.size()).cuda()))