def build_model(self):
data_shape = [None, 3, self.cfg.crop_size, self.cfg.crop_size]
input_A = fluid.data(name='input_A', shape=data_shape, dtype='float32')
input_B = fluid.data(name='input_B', shape=data_shape, dtype='float32')
fake_pool_A = fluid.data(
name='fake_pool_A', shape=data_shape, dtype='float32')
fake_pool_B = fluid.data(
name='fake_pool_B', shape=data_shape, dtype='float32')
# used for continuous evaluation
if self.cfg.enable_ce:
fluid.default_startup_program().random_seed = 90
A_py_reader = fluid.io.PyReader(
feed_list=[input_A],
capacity=4,
iterable=True,
use_double_buffer=True)
B_py_reader = fluid.io.PyReader(
feed_list=[input_B],
capacity=4,
iterable=True,
use_double_buffer=True)
gen_trainer = GTrainer(input_A, input_B, self.cfg, self.batch_num)
d_A_trainer = DATrainer(input_B, fake_pool_B, self.cfg, self.batch_num)
d_B_trainer = DBTrainer(input_A, fake_pool_A, self.cfg, self.batch_num)
# prepare environment
place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace()
A_py_reader.decorate_batch_generator(
self.A_reader,
places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places())
B_py_reader.decorate_batch_generator(
self.B_reader,
places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places())
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
A_pool = utility.ImagePool()
B_pool = utility.ImagePool()
if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, d_A_trainer, "net_DA")
utility.init_checkpoints(self.cfg, exe, d_B_trainer, "net_DB")
### memory optim
build_strategy = fluid.BuildStrategy()
build_strategy.enable_inplace = True
gen_trainer_program = fluid.CompiledProgram(
gen_trainer.program).with_data_parallel(
loss_name=gen_trainer.g_loss.name,
build_strategy=build_strategy)
d_A_trainer_program = fluid.CompiledProgram(
d_A_trainer.program).with_data_parallel(
loss_name=d_A_trainer.d_loss_A.name,
build_strategy=build_strategy)
d_B_trainer_program = fluid.CompiledProgram(
d_B_trainer.program).with_data_parallel(
loss_name=d_B_trainer.d_loss_B.name,
build_strategy=build_strategy)
t_time = 0
for epoch_id in range(self.cfg.epoch):
batch_id = 0
for data_A, data_B in zip(A_py_reader(), B_py_reader()):
s_time = time.time()
tensor_A, tensor_B = data_A[0]['input_A'], data_B[0]['input_B']
## optimize the g_A network
g_A_loss, g_A_cyc_loss, g_A_idt_loss, g_B_loss, g_B_cyc_loss,\
g_B_idt_loss, fake_A_tmp, fake_B_tmp = exe.run(
gen_trainer_program,
fetch_list=[
gen_trainer.G_A, gen_trainer.cyc_A_loss,
gen_trainer.idt_loss_A, gen_trainer.G_B,
gen_trainer.cyc_B_loss, gen_trainer.idt_loss_B,
gen_trainer.fake_A, gen_trainer.fake_B
],
feed={"input_A": tensor_A,
"input_B": tensor_B})
fake_pool_B = B_pool.pool_image(fake_B_tmp)
fake_pool_A = A_pool.pool_image(fake_A_tmp)
if self.cfg.enable_ce:
fake_pool_B = fake_B_tmp
fake_pool_A = fake_A_tmp
# optimize the d_A network
d_A_loss = exe.run(
d_A_trainer_program,
fetch_list=[d_A_trainer.d_loss_A],
feed={"input_B": tensor_B,
"fake_pool_B": fake_pool_B})[0]
# optimize the d_B network
d_B_loss = exe.run(
d_B_trainer_program,
fetch_list=[d_B_trainer.d_loss_B],
feed={"input_A": tensor_A,
"fake_pool_A": fake_pool_A})[0]
batch_time = time.time() - s_time
t_time += batch_time
if batch_id % self.cfg.print_freq == 0:
print("epoch{}: batch{}: \n\
d_A_loss: {}; g_A_loss: {}; g_A_cyc_loss: {}; g_A_idt_loss: {}; \n\
d_B_loss: {}; g_B_loss: {}; g_B_cyc_loss: {}; g_B_idt_loss: {}; \n\
Batch_time_cost: {}".format(
epoch_id, batch_id, d_A_loss[0], g_A_loss[0],
g_A_cyc_loss[0], g_A_idt_loss[0], d_B_loss[0], g_B_loss[
0], g_B_cyc_loss[0], g_B_idt_loss[0], batch_time))
sys.stdout.flush()
batch_id += 1
# used for continuous evaluation
if self.cfg.enable_ce and batch_id == 10:
break
if self.cfg.run_test:
A_image_name = fluid.data(
name='A_image_name', shape=[None, 1], dtype='int32')
B_image_name = fluid.data(
name='B_image_name', shape=[None, 1], dtype='int32')
A_test_py_reader = fluid.io.PyReader(
feed_list=[input_A, A_image_name],
capacity=4,
iterable=True,
use_double_buffer=True)
B_test_py_reader = fluid.io.PyReader(
feed_list=[input_B, B_image_name],
capacity=4,
iterable=True,
use_double_buffer=True)
A_test_py_reader.decorate_batch_generator(
self.A_test_reader,
places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places())
B_test_py_reader.decorate_batch_generator(
self.B_test_reader,
places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places())
test_program = gen_trainer.infer_program
utility.save_test_image(
epoch_id,
self.cfg,
exe,
place,
test_program,
gen_trainer,
A_test_py_reader,
B_test_py_reader,
A_id2name=self.A_id2name,
B_id2name=self.B_id2name)
if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer,
"net_G")
utility.checkpoints(epoch_id, self.cfg, exe, d_A_trainer,
"net_DA")
utility.checkpoints(epoch_id, self.cfg, exe, d_B_trainer,
"net_DB")
# used for continuous evaluation
if self.cfg.enable_ce:
device_num = fluid.core.get_cuda_device_count(
) if self.cfg.use_gpu else 1
print("kpis\tcyclegan_g_A_loss_card{}\t{}".format(device_num,
g_A_loss[0]))
print("kpis\tcyclegan_g_A_cyc_loss_card{}\t{}".format(
device_num, g_A_cyc_loss[0]))
print("kpis\tcyclegan_g_A_idt_loss_card{}\t{}".format(
device_num, g_A_idt_loss[0]))
print("kpis\tcyclegan_d_A_loss_card{}\t{}".format(device_num,
d_A_loss[0]))
print("kpis\tcyclegan_g_B_loss_card{}\t{}".format(device_num,
g_B_loss[0]))
print("kpis\tcyclegan_g_B_cyc_loss_card{}\t{}".format(
device_num, g_B_cyc_loss[0]))
print("kpis\tcyclegan_g_B_idt_loss_card{}\t{}".format(
device_num, g_B_idt_loss[0]))
print("kpis\tcyclegan_d_B_loss_card{}\t{}".format(device_num,
d_B_loss[0]))
print("kpis\tcyclegan_Batch_time_cost_card{}\t{}".format(
device_num, batch_time))
def build_model(self):
data_shape = [-1, 3, self.cfg.crop_size, self.cfg.crop_size]
input_A = fluid.layers.data(name='input_A',
shape=data_shape,
dtype='float32')
input_B = fluid.layers.data(name='input_B',
shape=data_shape,
dtype='float32')
fake_pool_A = fluid.layers.data(name='fake_pool_A',
shape=data_shape,
dtype='float32')
fake_pool_B = fluid.layers.data(name='fake_pool_B',
shape=data_shape,
dtype='float32')
gen_trainer = GTrainer(input_A, input_B, self.cfg, self.batch_num)
d_A_trainer = DATrainer(input_B, fake_pool_B, self.cfg, self.batch_num)
d_B_trainer = DBTrainer(input_A, fake_pool_A, self.cfg, self.batch_num)
# prepare environment
place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
A_pool = utility.ImagePool()
B_pool = utility.ImagePool()
if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, d_A_trainer, "net_DA")
utility.init_checkpoints(self.cfg, exe, d_B_trainer, "net_DB")
### memory optim
build_strategy = fluid.BuildStrategy()
build_strategy.enable_inplace = False
build_strategy.memory_optimize = False
gen_trainer_program = fluid.CompiledProgram(
gen_trainer.program).with_data_parallel(
loss_name=gen_trainer.g_loss.name,
build_strategy=build_strategy)
d_A_trainer_program = fluid.CompiledProgram(
d_A_trainer.program).with_data_parallel(
loss_name=d_A_trainer.d_loss_A.name,
build_strategy=build_strategy)
d_B_trainer_program = fluid.CompiledProgram(
d_B_trainer.program).with_data_parallel(
loss_name=d_B_trainer.d_loss_B.name,
build_strategy=build_strategy)
losses = [[], []]
t_time = 0
for epoch_id in range(self.cfg.epoch):
batch_id = 0
for i in range(self.batch_num):
data_A = next(self.A_reader())
data_B = next(self.B_reader())
tensor_A = fluid.LoDTensor()
tensor_B = fluid.LoDTensor()
tensor_A.set(data_A, place)
tensor_B.set(data_B, place)
s_time = time.time()
# optimize the g_A network
g_A_loss, g_A_cyc_loss, g_A_idt_loss, g_B_loss, g_B_cyc_loss,\
g_B_idt_loss, fake_A_tmp, fake_B_tmp = exe.run(
gen_trainer_program,
fetch_list=[
gen_trainer.G_A, gen_trainer.cyc_A_loss,
gen_trainer.idt_loss_A, gen_trainer.G_B,
gen_trainer.cyc_B_loss, gen_trainer.idt_loss_B,
gen_trainer.fake_A, gen_trainer.fake_B
],
feed={"input_A": tensor_A,
"input_B": tensor_B})
fake_pool_B = B_pool.pool_image(fake_B_tmp)
fake_pool_A = A_pool.pool_image(fake_A_tmp)
# optimize the d_A network
d_A_loss = exe.run(d_A_trainer_program,
fetch_list=[d_A_trainer.d_loss_A],
feed={
"input_B": tensor_B,
"fake_pool_B": fake_pool_B
})[0]
# optimize the d_B network
d_B_loss = exe.run(d_B_trainer_program,
fetch_list=[d_B_trainer.d_loss_B],
feed={
"input_A": tensor_A,
"fake_pool_A": fake_pool_A
})[0]
batch_time = time.time() - s_time
t_time += batch_time
if batch_id % self.cfg.print_freq == 0:
print("epoch{}: batch{}: \n\
d_A_loss: {}; g_A_loss: {}; g_A_cyc_loss: {}; g_A_idt_loss: {}; \n\
d_B_loss: {}; g_B_loss: {}; g_B_cyc_loss: {}; g_B_idt_loss: {}; \n\
Batch_time_cost: {:.2f}".format(
epoch_id, batch_id, d_A_loss[0], g_A_loss[0],
g_A_cyc_loss[0], g_A_idt_loss[0], d_B_loss[0],
g_B_loss[0], g_B_cyc_loss[0], g_B_idt_loss[0],
batch_time))
losses[0].append(g_A_loss[0])
losses[1].append(d_A_loss[0])
sys.stdout.flush()
batch_id += 1
if self.cfg.run_test:
test_program = gen_trainer.infer_program
utility.save_test_image(epoch_id, self.cfg, exe, place,
test_program, gen_trainer,
self.A_test_reader, self.B_test_reader)
if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer,
"net_G")
utility.checkpoints(epoch_id, self.cfg, exe, d_A_trainer,
"net_DA")
utility.checkpoints(epoch_id, self.cfg, exe, d_B_trainer,
"net_DB")