def create_model(n_res_blocks=9, device="cpu"):
"""Builds the generators and discriminators."""
# Instantiate generators
G_XtoY = CycleGenerator(n_residual_blocks=n_res_blocks)
G_YtoX = CycleGenerator(n_residual_blocks=n_res_blocks)
# Instantiate discriminators
D_X = Discriminator()
D_Y = Discriminator()
# move models to GPU, if available
if device != "cpu":
G_XtoY.to(device)
G_YtoX.to(device)
D_X.to(device)
D_Y.to(device)
print('Models moved to GPU.')
else:
print('Only CPU available.')
return G_XtoY, G_YtoX, D_X, D_Y
def main():
place = set_device(FLAGS.device)
fluid.enable_dygraph(place) if FLAGS.dynamic else None
# Generators
g_AB = Generator()
g_BA = Generator()
# Discriminators
d_A = Discriminator()
d_B = Discriminator()
g = GeneratorCombine(g_AB, g_BA, d_A, d_B)
da_params = d_A.parameters()
db_params = d_B.parameters()
g_params = g_AB.parameters() + g_BA.parameters()
da_optimizer = opt(da_params)
db_optimizer = opt(db_params)
g_optimizer = opt(g_params)
im_shape = [None, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
fake_A = Input(im_shape, 'float32', 'fake_A')
fake_B = Input(im_shape, 'float32', 'fake_B')
g_AB.prepare(inputs=[input_A], device=FLAGS.device)
g_BA.prepare(inputs=[input_B], device=FLAGS.device)
g.prepare(g_optimizer,
GLoss(),
inputs=[input_A, input_B],
device=FLAGS.device)
d_A.prepare(da_optimizer,
DLoss(),
inputs=[input_B, fake_B],
device=FLAGS.device)
d_B.prepare(db_optimizer,
DLoss(),
inputs=[input_A, fake_A],
device=FLAGS.device)
if FLAGS.resume:
g.load(FLAGS.resume)
loader_A = paddle.io.DataLoader(data.DataA(),
places=place,
shuffle=True,
return_list=True,
batch_size=FLAGS.batch_size)
loader_B = paddle.io.DataLoader(data.DataB(),
places=place,
shuffle=True,
return_list=True,
batch_size=FLAGS.batch_size)
A_pool = data.ImagePool()
B_pool = data.ImagePool()
for epoch in range(FLAGS.epoch):
for i, (data_A, data_B) in enumerate(zip(loader_A, loader_B)):
data_A = data_A[0][0] if not FLAGS.dynamic else data_A[0]
data_B = data_B[0][0] if not FLAGS.dynamic else data_B[0]
start = time.time()
fake_B = g_AB.test_batch(data_A)[0]
fake_A = g_BA.test_batch(data_B)[0]
g_loss = g.train_batch([data_A, data_B])[0]
fake_pb = B_pool.get(fake_B)
da_loss = d_A.train_batch([data_B, fake_pb])[0]
fake_pa = A_pool.get(fake_A)
db_loss = d_B.train_batch([data_A, fake_pa])[0]
t = time.time() - start
if i % 20 == 0:
print("epoch: {} | step: {:3d} | g_loss: {:.4f} | " \
"da_loss: {:.4f} | db_loss: {:.4f} | s/step {:.4f}".
format(epoch, i, g_loss[0], da_loss[0], db_loss[0], t))
g.save('{}/{}'.format(FLAGS.checkpoint_path, epoch))
default='/media/mint/Barracuda/Models/cyclegan/vangog',
help='path to store the models')
opt = parser.parse_args()
print(opt)
if torch.cuda.is_available() and not opt.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
###### Definition of variables ######
# Networks
netG_A2B = Generator(opt.input_nc, opt.output_nc)
netG_B2A = Generator(opt.output_nc, opt.input_nc)
netD_A = Discriminator(opt.input_nc)
netD_B = Discriminator(opt.output_nc)
if opt.cuda:
netG_A2B.cuda()
netG_B2A.cuda()
netD_A.cuda()
netD_B.cuda()
netG_A2B.apply(weights_init_normal)
netG_B2A.apply(weights_init_normal)
netD_A.apply(weights_init_normal)
netD_B.apply(weights_init_normal)
# Lossess
criterion_GAN = torch.nn.MSELoss()