未验证 提交 f1e2c268 编写于 作者: C ceci3 提交者: GitHub

Sync gan 1.7 (#4288)

* sync_1.7

* update save
上级 40ca043a
...@@ -23,7 +23,6 @@ from PIL import Image ...@@ -23,7 +23,6 @@ from PIL import Image
import paddle.fluid as fluid import paddle.fluid as fluid
import paddle import paddle
import numpy as np import numpy as np
import imageio
import glob import glob
from util.config import add_arguments, print_arguments from util.config import add_arguments, print_arguments
from data_reader import celeba_reader_creator, reader_creator, triplex_reader_creator from data_reader import celeba_reader_creator, reader_creator, triplex_reader_creator
...@@ -72,14 +71,14 @@ add_arg('no_instance', type=bool, default=False, help="Whether to use instance l ...@@ -72,14 +71,14 @@ add_arg('no_instance', type=bool, default=False, help="Whether to use instance l
def infer(args): def infer(args):
data_shape = [-1, 3, args.image_size, args.image_size] data_shape = [None, 3, args.image_size, args.image_size]
input = fluid.layers.data(name='input', shape=data_shape, dtype='float32') input = fluid.data(name='input', shape=data_shape, dtype='float32')
label_org_ = fluid.layers.data( label_org_ = fluid.data(
name='label_org_', shape=[args.c_dim], dtype='float32') name='label_org_', shape=[None, args.c_dim], dtype='float32')
label_trg_ = fluid.layers.data( label_trg_ = fluid.data(
name='label_trg_', shape=[args.c_dim], dtype='float32') name='label_trg_', shape=[None, args.c_dim], dtype='float32')
image_name = fluid.layers.data( image_name = fluid.data(
name='image_name', shape=[args.n_samples], dtype='int32') name='image_name', shape=[None, args.n_samples], dtype='int32')
model_name = 'net_G' model_name = 'net_G'
...@@ -110,7 +109,7 @@ def infer(args): ...@@ -110,7 +109,7 @@ def infer(args):
fake = model.network_G(input, "generator", cfg=args) fake = model.network_G(input, "generator", cfg=args)
elif args.model_net == 'StarGAN': elif args.model_net == 'StarGAN':
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[input, label_org_, label_trg_, image_name], feed_list=[input, label_org_, label_trg_, image_name],
capacity=32, capacity=32,
iterable=True, iterable=True,
...@@ -122,7 +121,7 @@ def infer(args): ...@@ -122,7 +121,7 @@ def infer(args):
elif args.model_net == 'STGAN': elif args.model_net == 'STGAN':
from network.STGAN_network import STGAN_model from network.STGAN_network import STGAN_model
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[input, label_org_, label_trg_, image_name], feed_list=[input, label_org_, label_trg_, image_name],
capacity=32, capacity=32,
iterable=True, iterable=True,
...@@ -139,7 +138,7 @@ def infer(args): ...@@ -139,7 +138,7 @@ def infer(args):
elif args.model_net == 'AttGAN': elif args.model_net == 'AttGAN':
from network.AttGAN_network import AttGAN_model from network.AttGAN_network import AttGAN_model
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[input, label_org_, label_trg_, image_name], feed_list=[input, label_org_, label_trg_, image_name],
capacity=32, capacity=32,
iterable=True, iterable=True,
...@@ -154,17 +153,17 @@ def infer(args): ...@@ -154,17 +153,17 @@ def infer(args):
name='generator', name='generator',
is_test=True) is_test=True)
elif args.model_net == 'CGAN': elif args.model_net == 'CGAN':
noise = fluid.layers.data( noise = fluid.data(
name='noise', shape=[args.noise_size], dtype='float32') name='noise', shape=[None, args.noise_size], dtype='float32')
conditions = fluid.layers.data( conditions = fluid.data(
name='conditions', shape=[1], dtype='float32') name='conditions', shape=[None, 1], dtype='float32')
from network.CGAN_network import CGAN_model from network.CGAN_network import CGAN_model
model = CGAN_model(args.n_samples) model = CGAN_model(args.n_samples)
fake = model.network_G(noise, conditions, name="G") fake = model.network_G(noise, conditions, name="G")
elif args.model_net == 'DCGAN': elif args.model_net == 'DCGAN':
noise = fluid.layers.data( noise = fluid.data(
name='noise', shape=[args.noise_size], dtype='float32') name='noise', shape=[None, args.noise_size], dtype='float32')
from network.DCGAN_network import DCGAN_model from network.DCGAN_network import DCGAN_model
model = DCGAN_model(args.n_samples) model = DCGAN_model(args.n_samples)
...@@ -197,10 +196,10 @@ def infer(args): ...@@ -197,10 +196,10 @@ def infer(args):
place = fluid.CUDAPlace(0) place = fluid.CUDAPlace(0)
exe = fluid.Executor(place) exe = fluid.Executor(place)
exe.run(fluid.default_startup_program()) exe.run(fluid.default_startup_program())
for var in fluid.default_main_program().global_block().all_parameters(): for var in fluid.default_main_program().all_parameters():
print(var.name) print(var.name)
print(args.init_model + '/' + model_name) print(args.init_model + '/' + model_name)
fluid.io.load_persistables(exe, os.path.join(args.init_model, model_name)) fluid.load(fluid.default_main_program(), os.path.join(args.init_model, model_name))
print('load params done') print('load params done')
if not os.path.exists(args.output): if not os.path.exists(args.output):
os.makedirs(args.output) os.makedirs(args.output)
...@@ -214,10 +213,10 @@ def infer(args): ...@@ -214,10 +213,10 @@ def infer(args):
args=args, args=args,
mode="VAL") mode="VAL")
reader_test = test_reader.make_reader(return_name=True) reader_test = test_reader.make_reader(return_name=True)
py_reader.decorate_batch_generator( loader.set_batch_generator(
reader_test, reader_test,
places=fluid.cuda_places() if args.use_gpu else fluid.cpu_places()) places=fluid.cuda_places() if args.use_gpu else fluid.cpu_places())
for data in py_reader(): for data in loader():
real_img, label_org, label_trg, image_name = data[0]['input'], data[ real_img, label_org, label_trg, image_name = data[0]['input'], data[
0]['label_org_'], data[0]['label_trg_'], data[0]['image_name'] 0]['label_org_'], data[0]['label_trg_'], data[0]['image_name']
image_name_save = _compute_start_end(image_name) image_name_save = _compute_start_end(image_name)
...@@ -251,9 +250,8 @@ def infer(args): ...@@ -251,9 +250,8 @@ def infer(args):
images_concat = np.concatenate(images, 1) images_concat = np.concatenate(images, 1)
if len(np.array(label_org)) > 1: if len(np.array(label_org)) > 1:
images_concat = np.concatenate(images_concat, 1) images_concat = np.concatenate(images_concat, 1)
imageio.imwrite( fake_image = Image.fromarray(((images_concat + 1) * 127.5).astype(np.uint8))
os.path.join(args.output, "fake_img_" + image_name_save), ( fake_image.save(os.path.join(args.output, "fake_image_" + image_name_save))
(images_concat + 1) * 127.5).astype(np.uint8))
elif args.model_net == 'StarGAN': elif args.model_net == 'StarGAN':
test_reader = celeba_reader_creator( test_reader = celeba_reader_creator(
image_dir=args.dataset_dir, image_dir=args.dataset_dir,
...@@ -261,10 +259,10 @@ def infer(args): ...@@ -261,10 +259,10 @@ def infer(args):
args=args, args=args,
mode="VAL") mode="VAL")
reader_test = test_reader.make_reader(return_name=True) reader_test = test_reader.make_reader(return_name=True)
py_reader.decorate_batch_generator( loader.set_batch_generator(
reader_test, reader_test,
places=fluid.cuda_places() if args.use_gpu else fluid.cpu_places()) places=fluid.cuda_places() if args.use_gpu else fluid.cpu_places())
for data in py_reader(): for data in loader():
real_img, label_org, label_trg, image_name = data[0]['input'], data[ real_img, label_org, label_trg, image_name = data[0]['input'], data[
0]['label_org_'], data[0]['label_trg_'], data[0]['image_name'] 0]['label_org_'], data[0]['label_trg_'], data[0]['image_name']
image_name_save = _compute_start_end(image_name) image_name_save = _compute_start_end(image_name)
...@@ -287,9 +285,8 @@ def infer(args): ...@@ -287,9 +285,8 @@ def infer(args):
images_concat = np.concatenate(images, 1) images_concat = np.concatenate(images, 1)
if len(np.array(label_org)) > 1: if len(np.array(label_org)) > 1:
images_concat = np.concatenate(images_concat, 1) images_concat = np.concatenate(images_concat, 1)
imageio.imwrite( fake_image = Image.fromarray(((images_concat + 1) * 127.5).astype(np.uint8))
os.path.join(args.output, "fake_img_" + image_name_save), ( fake_image.save(os.path.join(args.output, "fake_image_" + image_name_save))
(images_concat + 1) * 127.5).astype(np.uint8))
elif args.model_net == 'Pix2pix' or args.model_net == 'CycleGAN': elif args.model_net == 'Pix2pix' or args.model_net == 'CycleGAN':
test_reader = reader_creator( test_reader = reader_creator(
...@@ -315,12 +312,10 @@ def infer(args): ...@@ -315,12 +312,10 @@ def infer(args):
input_temp = save_batch_image(np.array(real_img)) input_temp = save_batch_image(np.array(real_img))
for i, name in enumerate(image_names): for i, name in enumerate(image_names):
imageio.imwrite( fake_image = Image.fromarray(((fake_temp[i] + 1) * 127.5).astype(np.uint8))
os.path.join(args.output, "fake_" + name), ( fake_image.save(os.path.join(args.output, "fake_" + name))
(fake_temp[i] + 1) * 127.5).astype(np.uint8)) input_image = Image.fromarray(((input_temp[i] + 1) * 127.5).astype(np.uint8))
imageio.imwrite( input_image.save(os.path.join(args.output, "input_" + name))
os.path.join(args.output, "input_" + name), (
(input_temp[i] + 1) * 127.5).astype(np.uint8))
elif args.model_net == 'SPADE': elif args.model_net == 'SPADE':
test_reader = triplex_reader_creator( test_reader = triplex_reader_creator(
image_dir=args.dataset_dir, image_dir=args.dataset_dir,
...@@ -345,10 +340,10 @@ def infer(args): ...@@ -345,10 +340,10 @@ def infer(args):
fake_B_temp = np.squeeze(fake_B_temp[0]).transpose([1, 2, 0]) fake_B_temp = np.squeeze(fake_B_temp[0]).transpose([1, 2, 0])
input_B_temp = np.squeeze(data_B[0]).transpose([1, 2, 0]) input_B_temp = np.squeeze(data_B[0]).transpose([1, 2, 0])
imageio.imwrite(args.output + "/fakeB_" + "_" + name, ( fakeB_image = Image.fromarray(((fake_B_temp + 1) * 127.5).astype(np.uint8))
(fake_B_temp + 1) * 127.5).astype(np.uint8)) fakeB_image.save(os.path.join(args.output, "fakeB_" + name))
imageio.imwrite(args.output + "/real_" + "_" + name, ( real_image = Image.fromarray(((input_B_temp + 1) * 127.5).astype(np.uint8))
(input_B_temp + 1) * 127.5).astype(np.uint8)) real_image.save(os.path.join(args.output, "real_" + name))
elif args.model_net == 'CGAN': elif args.model_net == 'CGAN':
noise_data = np.random.uniform( noise_data = np.random.uniform(
......
python infer.py --model_net STGAN --init_model ./output/stgan/checkpoints/19/ --dataset_dir ./data/celeba/ --image_size 128 --use_gru True --output ./infer_result/stgan/ python infer.py --model_net STGAN --init_model ./output/stgan/checkpoints/49/ --dataset_dir ./data/celeba/ --image_size 128 --use_gru True --output ./infer_result/stgan/
...@@ -31,8 +31,6 @@ class GTrainer(): ...@@ -31,8 +31,6 @@ class GTrainer():
model = AttGAN_model() model = AttGAN_model()
self.fake_img, self.rec_img = model.network_G( self.fake_img, self.rec_img = model.network_G(
image_real, label_org_, label_trg_, cfg, name="generator") image_real, label_org_, label_trg_, cfg, name="generator")
self.fake_img.persistable = True
self.rec_img.persistable = True
self.infer_program = self.program.clone(for_test=True) self.infer_program = self.program.clone(for_test=True)
self.g_loss_rec = fluid.layers.mean( self.g_loss_rec = fluid.layers.mean(
...@@ -64,9 +62,6 @@ class GTrainer(): ...@@ -64,9 +62,6 @@ class GTrainer():
label_trg)) label_trg))
self.g_loss = self.g_loss_fake + cfg.lambda_rec * self.g_loss_rec + cfg.lambda_cls * self.g_loss_cls self.g_loss = self.g_loss_fake + cfg.lambda_rec * self.g_loss_rec + cfg.lambda_cls * self.g_loss_cls
self.g_loss_fake.persistable = True
self.g_loss_rec.persistable = True
self.g_loss_cls.persistable = True
lr = fluid.layers.piecewise_decay( lr = fluid.layers.piecewise_decay(
boundaries=[99 * step_per_epoch], boundaries=[99 * step_per_epoch],
values=[cfg.g_lr, cfg.g_lr * 0.1]) values=[cfg.g_lr, cfg.g_lr * 0.1])
...@@ -133,11 +128,6 @@ class DTrainer(): ...@@ -133,11 +128,6 @@ class DTrainer():
raise NotImplementedError("gan_mode {} is not support!".format( raise NotImplementedError("gan_mode {} is not support!".format(
cfg.gan_mode)) cfg.gan_mode))
self.d_loss_real.persistable = True
self.d_loss_fake.persistable = True
self.d_loss.persistable = True
self.d_loss_cls.persistable = True
self.d_loss_gp.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
if fluid.io.is_parameter(var) and var.name.startswith( if fluid.io.is_parameter(var) and var.name.startswith(
...@@ -300,7 +290,7 @@ class AttGAN(object): ...@@ -300,7 +290,7 @@ class AttGAN(object):
if self.cfg.enable_ce: if self.cfg.enable_ce:
fluid.default_startup_program().random_seed = 90 fluid.default_startup_program().random_seed = 90
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[image_real, label_org, label_trg], feed_list=[image_real, label_org, label_trg],
capacity=64, capacity=64,
iterable=True, iterable=True,
...@@ -320,7 +310,7 @@ class AttGAN(object): ...@@ -320,7 +310,7 @@ class AttGAN(object):
# prepare environment # prepare environment
place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace() place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace()
py_reader.decorate_batch_generator( loader.set_batch_generator(
self.train_reader, self.train_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -328,8 +318,8 @@ class AttGAN(object): ...@@ -328,8 +318,8 @@ class AttGAN(object):
exe.run(fluid.default_startup_program()) exe.run(fluid.default_startup_program())
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G") utility.init_checkpoints(self.cfg, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, dis_trainer, "net_D") utility.init_checkpoints(self.cfg, dis_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -351,7 +341,7 @@ class AttGAN(object): ...@@ -351,7 +341,7 @@ class AttGAN(object):
for epoch_id in range(self.cfg.epoch): for epoch_id in range(self.cfg.epoch):
batch_id = 0 batch_id = 0
for data in py_reader(): for data in loader():
s_time = time.time() s_time = time.time()
# optimize the discriminator network # optimize the discriminator network
fetches = [ fetches = [
...@@ -396,12 +386,12 @@ class AttGAN(object): ...@@ -396,12 +386,12 @@ class AttGAN(object):
name='image_name', name='image_name',
shape=[None, self.cfg.n_samples], shape=[None, self.cfg.n_samples],
dtype='int32') dtype='int32')
test_py_reader = fluid.io.PyReader( test_loader = fluid.io.DataLoader.from_generator(
feed_list=[image_real, label_org, label_trg, image_name], feed_list=[image_real, label_org, label_trg, image_name],
capacity=32, capacity=32,
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
test_py_reader.decorate_batch_generator( test_loader.set_batch_generator(
self.test_reader, self.test_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -409,12 +399,12 @@ class AttGAN(object): ...@@ -409,12 +399,12 @@ class AttGAN(object):
test_program = test_gen_trainer.infer_program test_program = test_gen_trainer.infer_program
utility.save_test_image(epoch_id, self.cfg, exe, place, utility.save_test_image(epoch_id, self.cfg, exe, place,
test_program, test_gen_trainer, test_program, test_gen_trainer,
test_py_reader) test_loader)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, utility.checkpoints(epoch_id, self.cfg, gen_trainer,
"net_G") "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, dis_trainer, utility.checkpoints(epoch_id, self.cfg, dis_trainer,
"net_D") "net_D")
# used for continuous evaluation # used for continuous evaluation
if self.cfg.enable_ce: if self.cfg.enable_ce:
......
...@@ -35,7 +35,6 @@ class GTrainer(): ...@@ -35,7 +35,6 @@ class GTrainer():
with fluid.program_guard(self.program): with fluid.program_guard(self.program):
model = CGAN_model(cfg.batch_size) model = CGAN_model(cfg.batch_size)
self.fake = model.network_G(input, conditions, name="G") self.fake = model.network_G(input, conditions, name="G")
self.fake.persistable = True
self.infer_program = self.program.clone(for_test=True) self.infer_program = self.program.clone(for_test=True)
d_fake = model.network_D(self.fake, conditions, name="D") d_fake = model.network_D(self.fake, conditions, name="D")
fake_labels = fluid.layers.fill_constant_batch_size_like( fake_labels = fluid.layers.fill_constant_batch_size_like(
...@@ -43,7 +42,6 @@ class GTrainer(): ...@@ -43,7 +42,6 @@ class GTrainer():
self.g_loss = fluid.layers.reduce_mean( self.g_loss = fluid.layers.reduce_mean(
fluid.layers.sigmoid_cross_entropy_with_logits( fluid.layers.sigmoid_cross_entropy_with_logits(
x=d_fake, label=fake_labels)) x=d_fake, label=fake_labels))
self.g_loss.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
...@@ -64,7 +62,6 @@ class DTrainer(): ...@@ -64,7 +62,6 @@ class DTrainer():
self.d_loss = fluid.layers.reduce_mean( self.d_loss = fluid.layers.reduce_mean(
fluid.layers.sigmoid_cross_entropy_with_logits( fluid.layers.sigmoid_cross_entropy_with_logits(
x=d_logit, label=labels)) x=d_logit, label=labels))
self.d_loss.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
if fluid.io.is_parameter(var) and (var.name.startswith("D")): if fluid.io.is_parameter(var) and (var.name.startswith("D")):
...@@ -108,8 +105,8 @@ class CGAN(object): ...@@ -108,8 +105,8 @@ class CGAN(object):
size=[self.cfg.batch_size, self.cfg.noise_size]).astype('float32') size=[self.cfg.batch_size, self.cfg.noise_size]).astype('float32')
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, g_trainer, "net_G") utility.init_checkpoints(self.cfg, g_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, d_trainer, "net_D") utility.init_checkpoints(self.cfg, d_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -205,5 +202,5 @@ class CGAN(object): ...@@ -205,5 +202,5 @@ class CGAN(object):
plt.close(fig) plt.close(fig)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, g_trainer, "net_G") utility.checkpoints(epoch_id, self.cfg, g_trainer, "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, d_trainer, "net_D") utility.checkpoints(epoch_id, self.cfg, d_trainer, "net_D")
...@@ -33,9 +33,7 @@ class GTrainer(): ...@@ -33,9 +33,7 @@ class GTrainer():
with fluid.program_guard(self.program): with fluid.program_guard(self.program):
model = CycleGAN_model() model = CycleGAN_model()
self.fake_B = model.network_G(input_A, name="GA", cfg=cfg) self.fake_B = model.network_G(input_A, name="GA", cfg=cfg)
self.fake_B.persistable = True
self.fake_A = model.network_G(input_B, name="GB", cfg=cfg) self.fake_A = model.network_G(input_B, name="GB", cfg=cfg)
self.fake_A.persistable = True
self.cyc_A = model.network_G(self.fake_B, name="GB", cfg=cfg) self.cyc_A = model.network_G(self.fake_B, name="GB", cfg=cfg)
self.cyc_B = model.network_G(self.fake_A, name="GA", cfg=cfg) self.cyc_B = model.network_G(self.fake_A, name="GA", cfg=cfg)
...@@ -48,20 +46,16 @@ class GTrainer(): ...@@ -48,20 +46,16 @@ class GTrainer():
fluid.layers.elementwise_sub( fluid.layers.elementwise_sub(
x=input_B, y=self.cyc_B)) x=input_B, y=self.cyc_B))
self.cyc_A_loss = fluid.layers.reduce_mean(diff_A) * lambda_A self.cyc_A_loss = fluid.layers.reduce_mean(diff_A) * lambda_A
self.cyc_A_loss.persistable = True
self.cyc_B_loss = fluid.layers.reduce_mean(diff_B) * lambda_B self.cyc_B_loss = fluid.layers.reduce_mean(diff_B) * lambda_B
self.cyc_B_loss.persistable = True
self.cyc_loss = self.cyc_A_loss + self.cyc_B_loss self.cyc_loss = self.cyc_A_loss + self.cyc_B_loss
# GAN Loss D_A(G_A(A)) # GAN Loss D_A(G_A(A))
self.fake_rec_A = model.network_D(self.fake_B, name="DA", cfg=cfg) self.fake_rec_A = model.network_D(self.fake_B, name="DA", cfg=cfg)
self.G_A = fluid.layers.reduce_mean( self.G_A = fluid.layers.reduce_mean(
fluid.layers.square(self.fake_rec_A - 1)) fluid.layers.square(self.fake_rec_A - 1))
self.G_A.persistable = True
# GAN Loss D_B(G_B(B)) # GAN Loss D_B(G_B(B))
self.fake_rec_B = model.network_D(self.fake_A, name="DB", cfg=cfg) self.fake_rec_B = model.network_D(self.fake_A, name="DB", cfg=cfg)
self.G_B = fluid.layers.reduce_mean( self.G_B = fluid.layers.reduce_mean(
fluid.layers.square(self.fake_rec_B - 1)) fluid.layers.square(self.fake_rec_B - 1))
self.G_B.persistable = True
self.G = self.G_A + self.G_B self.G = self.G_A + self.G_B
# Identity Loss G_A # Identity Loss G_A
self.idt_A = model.network_G(input_B, name="GA", cfg=cfg) self.idt_A = model.network_G(input_B, name="GA", cfg=cfg)
...@@ -69,14 +63,12 @@ class GTrainer(): ...@@ -69,14 +63,12 @@ class GTrainer():
fluid.layers.abs( fluid.layers.abs(
fluid.layers.elementwise_sub( fluid.layers.elementwise_sub(
x=input_B, y=self.idt_A))) * lambda_B * lambda_identity x=input_B, y=self.idt_A))) * lambda_B * lambda_identity
self.idt_loss_A.persistable = True
# Identity Loss G_B # Identity Loss G_B
self.idt_B = model.network_G(input_A, name="GB", cfg=cfg) self.idt_B = model.network_G(input_A, name="GB", cfg=cfg)
self.idt_loss_B = fluid.layers.reduce_mean( self.idt_loss_B = fluid.layers.reduce_mean(
fluid.layers.abs( fluid.layers.abs(
fluid.layers.elementwise_sub( fluid.layers.elementwise_sub(
x=input_A, y=self.idt_B))) * lambda_A * lambda_identity x=input_A, y=self.idt_B))) * lambda_A * lambda_identity
self.idt_loss_B.persistable = True
self.idt_loss = fluid.layers.elementwise_add(self.idt_loss_A, self.idt_loss = fluid.layers.elementwise_add(self.idt_loss_A,
self.idt_loss_B) self.idt_loss_B)
...@@ -120,7 +112,6 @@ class DATrainer(): ...@@ -120,7 +112,6 @@ class DATrainer():
self.d_loss_A = (fluid.layers.square(self.fake_pool_rec_B) + self.d_loss_A = (fluid.layers.square(self.fake_pool_rec_B) +
fluid.layers.square(self.rec_B - 1)) / 2.0 fluid.layers.square(self.rec_B - 1)) / 2.0
self.d_loss_A = fluid.layers.reduce_mean(self.d_loss_A) self.d_loss_A = fluid.layers.reduce_mean(self.d_loss_A)
self.d_loss_A.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
...@@ -161,7 +152,6 @@ class DBTrainer(): ...@@ -161,7 +152,6 @@ class DBTrainer():
self.d_loss_B = (fluid.layers.square(self.fake_pool_rec_A) + self.d_loss_B = (fluid.layers.square(self.fake_pool_rec_A) +
fluid.layers.square(self.rec_A - 1)) / 2.0 fluid.layers.square(self.rec_A - 1)) / 2.0
self.d_loss_B = fluid.layers.reduce_mean(self.d_loss_B) self.d_loss_B = fluid.layers.reduce_mean(self.d_loss_B)
self.d_loss_B.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
if fluid.io.is_parameter(var) and var.name.startswith("DB"): if fluid.io.is_parameter(var) and var.name.startswith("DB"):
...@@ -244,13 +234,13 @@ class CycleGAN(object): ...@@ -244,13 +234,13 @@ class CycleGAN(object):
if self.cfg.enable_ce: if self.cfg.enable_ce:
fluid.default_startup_program().random_seed = 90 fluid.default_startup_program().random_seed = 90
A_py_reader = fluid.io.PyReader( A_loader = fluid.io.DataLoader.from_generator(
feed_list=[input_A], feed_list=[input_A],
capacity=4, capacity=4,
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
B_py_reader = fluid.io.PyReader( B_loader = fluid.io.DataLoader.from_generator(
feed_list=[input_B], feed_list=[input_B],
capacity=4, capacity=4,
iterable=True, iterable=True,
...@@ -263,11 +253,11 @@ class CycleGAN(object): ...@@ -263,11 +253,11 @@ class CycleGAN(object):
# prepare environment # prepare environment
place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace() place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace()
A_py_reader.decorate_batch_generator( A_loader.set_batch_generator(
self.A_reader, self.A_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
B_py_reader.decorate_batch_generator( B_loader.set_batch_generator(
self.B_reader, self.B_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -279,9 +269,9 @@ class CycleGAN(object): ...@@ -279,9 +269,9 @@ class CycleGAN(object):
B_pool = utility.ImagePool() B_pool = utility.ImagePool()
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G") utility.init_checkpoints(self.cfg, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, d_A_trainer, "net_DA") utility.init_checkpoints(self.cfg, d_A_trainer, "net_DA")
utility.init_checkpoints(self.cfg, exe, d_B_trainer, "net_DB") utility.init_checkpoints(self.cfg, d_B_trainer, "net_DB")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -304,7 +294,7 @@ class CycleGAN(object): ...@@ -304,7 +294,7 @@ class CycleGAN(object):
for epoch_id in range(self.cfg.epoch): for epoch_id in range(self.cfg.epoch):
batch_id = 0 batch_id = 0
for data_A, data_B in zip(A_py_reader(), B_py_reader()): for data_A, data_B in zip(A_loader(), B_loader()):
s_time = time.time() s_time = time.time()
tensor_A, tensor_B = data_A[0]['input_A'], data_B[0]['input_B'] tensor_A, tensor_B = data_A[0]['input_A'], data_B[0]['input_B']
## optimize the g_A network ## optimize the g_A network
...@@ -363,23 +353,23 @@ class CycleGAN(object): ...@@ -363,23 +353,23 @@ class CycleGAN(object):
name='A_image_name', shape=[None, 1], dtype='int32') name='A_image_name', shape=[None, 1], dtype='int32')
B_image_name = fluid.data( B_image_name = fluid.data(
name='B_image_name', shape=[None, 1], dtype='int32') name='B_image_name', shape=[None, 1], dtype='int32')
A_test_py_reader = fluid.io.PyReader( A_test_loader = fluid.io.DataLoader.from_generator(
feed_list=[input_A, A_image_name], feed_list=[input_A, A_image_name],
capacity=4, capacity=4,
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
B_test_py_reader = fluid.io.PyReader( B_test_loader = fluid.io.DataLoader.from_generator(
feed_list=[input_B, B_image_name], feed_list=[input_B, B_image_name],
capacity=4, capacity=4,
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
A_test_py_reader.decorate_batch_generator( A_test_loader.set_batch_generator(
self.A_test_reader, self.A_test_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
B_test_py_reader.decorate_batch_generator( B_test_loader.set_batch_generator(
self.B_test_reader, self.B_test_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -391,17 +381,17 @@ class CycleGAN(object): ...@@ -391,17 +381,17 @@ class CycleGAN(object):
place, place,
test_program, test_program,
gen_trainer, gen_trainer,
A_test_py_reader, A_test_loader,
B_test_py_reader, B_test_loader,
A_id2name=self.A_id2name, A_id2name=self.A_id2name,
B_id2name=self.B_id2name) B_id2name=self.B_id2name)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, utility.checkpoints(epoch_id, self.cfg, gen_trainer,
"net_G") "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, d_A_trainer, utility.checkpoints(epoch_id, self.cfg, d_A_trainer,
"net_DA") "net_DA")
utility.checkpoints(epoch_id, self.cfg, exe, d_B_trainer, utility.checkpoints(epoch_id, self.cfg, d_B_trainer,
"net_DB") "net_DB")
# used for continuous evaluation # used for continuous evaluation
......
...@@ -36,7 +36,6 @@ class GTrainer(): ...@@ -36,7 +36,6 @@ class GTrainer():
with fluid.program_guard(self.program): with fluid.program_guard(self.program):
model = DCGAN_model(cfg.batch_size) model = DCGAN_model(cfg.batch_size)
self.fake = model.network_G(input, name='G') self.fake = model.network_G(input, name='G')
self.fake.persistable = True
self.infer_program = self.program.clone(for_test=True) self.infer_program = self.program.clone(for_test=True)
d_fake = model.network_D(self.fake, name="D") d_fake = model.network_D(self.fake, name="D")
fake_labels = fluid.layers.fill_constant_batch_size_like( fake_labels = fluid.layers.fill_constant_batch_size_like(
...@@ -44,7 +43,6 @@ class GTrainer(): ...@@ -44,7 +43,6 @@ class GTrainer():
self.g_loss = fluid.layers.reduce_mean( self.g_loss = fluid.layers.reduce_mean(
fluid.layers.sigmoid_cross_entropy_with_logits( fluid.layers.sigmoid_cross_entropy_with_logits(
x=d_fake, label=fake_labels)) x=d_fake, label=fake_labels))
self.g_loss.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
...@@ -64,7 +62,6 @@ class DTrainer(): ...@@ -64,7 +62,6 @@ class DTrainer():
self.d_loss = fluid.layers.reduce_mean( self.d_loss = fluid.layers.reduce_mean(
fluid.layers.sigmoid_cross_entropy_with_logits( fluid.layers.sigmoid_cross_entropy_with_logits(
x=d_logit, label=labels)) x=d_logit, label=labels))
self.d_loss.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
if fluid.io.is_parameter(var) and (var.name.startswith("D")): if fluid.io.is_parameter(var) and (var.name.startswith("D")):
...@@ -113,8 +110,8 @@ class DCGAN(object): ...@@ -113,8 +110,8 @@ class DCGAN(object):
size=[self.cfg.batch_size, self.cfg.noise_size]).astype('float32') size=[self.cfg.batch_size, self.cfg.noise_size]).astype('float32')
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, g_trainer, "net_G") utility.init_checkpoints(self.cfg, g_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, d_trainer, "net_D") utility.init_checkpoints(self.cfg, d_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -207,8 +204,8 @@ class DCGAN(object): ...@@ -207,8 +204,8 @@ class DCGAN(object):
plt.close(fig) plt.close(fig)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, g_trainer, "net_G") utility.checkpoints(epoch_id, self.cfg, g_trainer, "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, d_trainer, "net_D") utility.checkpoints(epoch_id, self.cfg, d_trainer, "net_D")
# used for continuous evaluation # used for continuous evaluation
if self.cfg.enable_ce: if self.cfg.enable_ce:
device_num = fluid.core.get_cuda_device_count( device_num = fluid.core.get_cuda_device_count(
......
...@@ -30,7 +30,6 @@ class GTrainer(): ...@@ -30,7 +30,6 @@ class GTrainer():
with fluid.program_guard(self.program): with fluid.program_guard(self.program):
model = Pix2pix_model() model = Pix2pix_model()
self.fake_B = model.network_G(input_A, "generator", cfg=cfg) self.fake_B = model.network_G(input_A, "generator", cfg=cfg)
self.fake_B.persistable = True
self.infer_program = self.program.clone() self.infer_program = self.program.clone()
AB = fluid.layers.concat([input_A, self.fake_B], 1) AB = fluid.layers.concat([input_A, self.fake_B], 1)
self.pred = model.network_D(AB, "discriminator", cfg) self.pred = model.network_D(AB, "discriminator", cfg)
...@@ -246,8 +245,8 @@ class Pix2pix(object): ...@@ -246,8 +245,8 @@ class Pix2pix(object):
exe.run(fluid.default_startup_program()) exe.run(fluid.default_startup_program())
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G") utility.init_checkpoints(self.cfg, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, dis_trainer, "net_D") utility.init_checkpoints(self.cfg, dis_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -339,9 +338,9 @@ class Pix2pix(object): ...@@ -339,9 +338,9 @@ class Pix2pix(object):
A_id2name=self.id2name) A_id2name=self.id2name)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, utility.checkpoints(epoch_id, self.cfg, gen_trainer,
"net_G") "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, dis_trainer, utility.checkpoints(epoch_id, self.cfg, dis_trainer,
"net_D") "net_D")
if self.cfg.enable_ce: if self.cfg.enable_ce:
device_num = fluid.core.get_cuda_device_count( device_num = fluid.core.get_cuda_device_count(
......
...@@ -36,7 +36,6 @@ class GTrainer(): ...@@ -36,7 +36,6 @@ class GTrainer():
if not cfg.no_instance: if not cfg.no_instance:
input = fluid.layers.concat([input_label, input_ins], 1) input = fluid.layers.concat([input_label, input_ins], 1)
self.fake_B = model.network_G(input, "generator", cfg=cfg) self.fake_B = model.network_G(input, "generator", cfg=cfg)
self.fake_B.persistable = True
self.infer_program = self.program.clone() self.infer_program = self.program.clone()
fake_concat = fluid.layers.concat([input, self.fake_B], 1) fake_concat = fluid.layers.concat([input, self.fake_B], 1)
real_concat = fluid.layers.concat([input, input_img], 1) real_concat = fluid.layers.concat([input, input_img], 1)
...@@ -65,7 +64,6 @@ class GTrainer(): ...@@ -65,7 +64,6 @@ class GTrainer():
self.gan_loss /= len(self.pred_fake) self.gan_loss /= len(self.pred_fake)
else: else:
self.gan_loss = -1 * fluid.layers.reduce_mean(self.pred_fake) self.gan_loss = -1 * fluid.layers.reduce_mean(self.pred_fake)
self.gan_loss.persistable = True
#####GAN Feat loss #####GAN Feat loss
num_D = len(self.pred_fake) num_D = len(self.pred_fake)
self.gan_feat_loss = 0.0 self.gan_feat_loss = 0.0
...@@ -77,7 +75,6 @@ class GTrainer(): ...@@ -77,7 +75,6 @@ class GTrainer():
fluid.layers.elementwise_sub( fluid.layers.elementwise_sub(
x=self.pred_fake[i][j], y=self.pred_real[i][ x=self.pred_fake[i][j], y=self.pred_real[i][
j]))) * cfg.lambda_feat / num_D j]))) * cfg.lambda_feat / num_D
self.gan_feat_loss.persistable = True
########VGG Feat loss ########VGG Feat loss
weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0] weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
self.vgg = vgg.VGG19() self.vgg = vgg.VGG19()
...@@ -89,7 +86,6 @@ class GTrainer(): ...@@ -89,7 +86,6 @@ class GTrainer():
fluid.layers.abs( fluid.layers.abs(
fluid.layers.elementwise_sub( fluid.layers.elementwise_sub(
x=fake_vgg[i], y=real_vgg[i]))) x=fake_vgg[i], y=real_vgg[i])))
self.vgg_loss.persistable = True
self.g_loss = ( self.g_loss = (
self.gan_loss + self.gan_feat_loss + self.vgg_loss) / 3 self.gan_loss + self.gan_feat_loss + self.vgg_loss) / 3
lr = cfg.learning_rate lr = cfg.learning_rate
...@@ -173,8 +169,6 @@ class DTrainer(): ...@@ -173,8 +169,6 @@ class DTrainer():
loss_i = -1 * fluid.layers.reduce_mean(minval) loss_i = -1 * fluid.layers.reduce_mean(minval)
self.gan_loss_real += loss_i self.gan_loss_real += loss_i
self.gan_loss_real /= len(self.pred_real) self.gan_loss_real /= len(self.pred_real)
self.gan_loss_real.persistable = True
self.gan_loss_fake.persistable = True
self.d_loss = 0.5 * (self.gan_loss_real + self.gan_loss_fake) self.d_loss = 0.5 * (self.gan_loss_real + self.gan_loss_fake)
vars = [] vars = []
...@@ -310,12 +304,12 @@ class SPADE(object): ...@@ -310,12 +304,12 @@ class SPADE(object):
self.batch_num) self.batch_num)
dis_trainer = DTrainer(input_A, input_B, input_C, input_fake, self.cfg, dis_trainer = DTrainer(input_A, input_B, input_C, input_fake, self.cfg,
self.batch_num) self.batch_num)
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[input_A, input_B, input_C], feed_list=[input_A, input_B, input_C],
capacity=4, ## batch_size * 4 capacity=4, ## batch_size * 4
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
py_reader.decorate_batch_generator( loader.set_batch_generator(
self.train_reader, self.train_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -334,8 +328,8 @@ class SPADE(object): ...@@ -334,8 +328,8 @@ class SPADE(object):
self.cfg.vgg19_pretrain) self.cfg.vgg19_pretrain)
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G") utility.init_checkpoints(self.cfg, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, dis_trainer, "net_D") utility.init_checkpoints(self.cfg, dis_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -359,7 +353,7 @@ class SPADE(object): ...@@ -359,7 +353,7 @@ class SPADE(object):
for epoch_id in range(self.cfg.epoch): for epoch_id in range(self.cfg.epoch):
batch_id = 0 batch_id = 0
for tensor in py_reader(): for tensor in loader():
data_A, data_B, data_C = tensor[0]['input_label'], tensor[0][ data_A, data_B, data_C = tensor[0]['input_label'], tensor[0][
'input_img'], tensor[0]['input_ins'] 'input_img'], tensor[0]['input_ins']
s_time = time.time() s_time = time.time()
...@@ -408,12 +402,12 @@ class SPADE(object): ...@@ -408,12 +402,12 @@ class SPADE(object):
name='image_name', name='image_name',
shape=[None, self.cfg.batch_size], shape=[None, self.cfg.batch_size],
dtype="int32") dtype="int32")
test_py_reader = fluid.io.PyReader( test_loader = fluid.io.DataLoader.from_generator(
feed_list=[input_A, input_B, input_C, image_name], feed_list=[input_A, input_B, input_C, image_name],
capacity=4, ## batch_size * 4 capacity=4, ## batch_size * 4
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
test_py_reader.decorate_batch_generator( test_loader.set_batch_generator(
self.test_reader, self.test_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -424,13 +418,13 @@ class SPADE(object): ...@@ -424,13 +418,13 @@ class SPADE(object):
place, place,
test_program, test_program,
gen_trainer, gen_trainer,
test_py_reader, test_loader,
A_id2name=self.id2name) A_id2name=self.id2name)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, utility.checkpoints(epoch_id, self.cfg, gen_trainer,
"net_G") "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, dis_trainer, utility.checkpoints(epoch_id, self.cfg, dis_trainer,
"net_D") "net_D")
# used for continuous evaluation # used for continuous evaluation
if self.cfg.enable_ce: if self.cfg.enable_ce:
......
...@@ -33,8 +33,6 @@ class GTrainer(): ...@@ -33,8 +33,6 @@ class GTrainer():
model = STGAN_model() model = STGAN_model()
self.fake_img, self.rec_img = model.network_G( self.fake_img, self.rec_img = model.network_G(
image_real, label_org_, label_trg_, cfg, name="generator") image_real, label_org_, label_trg_, cfg, name="generator")
self.fake_img.persistable = True
self.rec_img.persistable = True
self.infer_program = self.program.clone(for_test=True) self.infer_program = self.program.clone(for_test=True)
self.g_loss_rec = fluid.layers.mean( self.g_loss_rec = fluid.layers.mean(
fluid.layers.abs( fluid.layers.abs(
...@@ -64,9 +62,6 @@ class GTrainer(): ...@@ -64,9 +62,6 @@ class GTrainer():
fluid.layers.sigmoid_cross_entropy_with_logits(self.cls_fake, fluid.layers.sigmoid_cross_entropy_with_logits(self.cls_fake,
label_trg)) label_trg))
self.g_loss = self.g_loss_fake + cfg.lambda_rec * self.g_loss_rec + cfg.lambda_cls * self.g_loss_cls self.g_loss = self.g_loss_fake + cfg.lambda_rec * self.g_loss_rec + cfg.lambda_cls * self.g_loss_cls
self.g_loss_fake.persistable = True
self.g_loss_rec.persistable = True
self.g_loss_cls.persistable = True
lr = cfg.g_lr lr = cfg.g_lr
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
...@@ -95,8 +90,6 @@ class DTrainer(): ...@@ -95,8 +90,6 @@ class DTrainer():
image_real, label_org_, label_trg_, cfg, name="generator") image_real, label_org_, label_trg_, cfg, name="generator")
self.pred_real, self.cls_real = model.network_D( self.pred_real, self.cls_real = model.network_D(
image_real, cfg, name="discriminator") image_real, cfg, name="discriminator")
self.pred_real.persistable = True
self.cls_real.persistable = True
self.pred_fake, _ = model.network_D( self.pred_fake, _ = model.network_D(
self.fake_img, cfg, name="discriminator") self.fake_img, cfg, name="discriminator")
self.d_loss_cls = fluid.layers.mean( self.d_loss_cls = fluid.layers.mean(
...@@ -137,11 +130,6 @@ class DTrainer(): ...@@ -137,11 +130,6 @@ class DTrainer():
raise NotImplementedError("gan_mode {} is not support!".format( raise NotImplementedError("gan_mode {} is not support!".format(
cfg.gan_mode)) cfg.gan_mode))
self.d_loss_real.persistable = True
self.d_loss_fake.persistable = True
self.d_loss.persistable = True
self.d_loss_cls.persistable = True
self.d_loss_gp.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
if fluid.io.is_parameter(var) and ( if fluid.io.is_parameter(var) and (
...@@ -158,8 +146,6 @@ class DTrainer(): ...@@ -158,8 +146,6 @@ class DTrainer():
name="net_D") name="net_D")
optimizer.minimize(self.d_loss, parameter_list=vars) optimizer.minimize(self.d_loss, parameter_list=vars)
f = open('G_program.txt', 'w')
print(self.program, file=f)
def gradient_penalty(self, f, real, fake=None, cfg=None, name=None): def gradient_penalty(self, f, real, fake=None, cfg=None, name=None):
def _interpolate(a, b=None): def _interpolate(a, b=None):
...@@ -317,7 +303,7 @@ class STGAN(object): ...@@ -317,7 +303,7 @@ class STGAN(object):
label_trg, label_trg_, self.cfg, label_trg, label_trg_, self.cfg,
self.batch_num) self.batch_num)
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[image_real, label_org, label_trg], feed_list=[image_real, label_org, label_trg],
capacity=64, capacity=64,
iterable=True, iterable=True,
...@@ -332,7 +318,7 @@ class STGAN(object): ...@@ -332,7 +318,7 @@ class STGAN(object):
# prepare environment # prepare environment
place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace() place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace()
py_reader.decorate_batch_generator( loader.set_batch_generator(
self.train_reader, self.train_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -341,8 +327,8 @@ class STGAN(object): ...@@ -341,8 +327,8 @@ class STGAN(object):
exe.run(fluid.default_startup_program()) exe.run(fluid.default_startup_program())
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G") utility.init_checkpoints(self.cfg, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, dis_trainer, "net_D") utility.init_checkpoints(self.cfg, dis_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -366,7 +352,7 @@ class STGAN(object): ...@@ -366,7 +352,7 @@ class STGAN(object):
for epoch_id in range(self.cfg.epoch): for epoch_id in range(self.cfg.epoch):
batch_id = 0 batch_id = 0
for data in py_reader(): for data in loader():
if self.cfg.max_iter and total_train_batch == self.cfg.max_iter: # used for benchmark if self.cfg.max_iter and total_train_batch == self.cfg.max_iter: # used for benchmark
return return
s_time = time.time() s_time = time.time()
...@@ -417,24 +403,24 @@ class STGAN(object): ...@@ -417,24 +403,24 @@ class STGAN(object):
name='image_name', name='image_name',
shape=[None, self.cfg.n_samples], shape=[None, self.cfg.n_samples],
dtype='int32') dtype='int32')
test_py_reader = fluid.io.PyReader( test_loader = fluid.io.DataLoader.from_generator(
feed_list=[image_real, label_org, label_trg, image_name], feed_list=[image_real, label_org, label_trg, image_name],
capacity=32, capacity=32,
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
test_py_reader.decorate_batch_generator( test_loader.set_batch_generator(
self.test_reader, self.test_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
test_program = test_gen_trainer.infer_program test_program = test_gen_trainer.infer_program
utility.save_test_image(epoch_id, self.cfg, exe, place, utility.save_test_image(epoch_id, self.cfg, exe, place,
test_program, test_gen_trainer, test_program, test_gen_trainer,
test_py_reader) test_loader)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, utility.checkpoints(epoch_id, self.cfg, gen_trainer,
"net_G") "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, dis_trainer, utility.checkpoints(epoch_id, self.cfg, dis_trainer,
"net_D") "net_D")
# used for continuous evaluation # used for continuous evaluation
if self.cfg.enable_ce: if self.cfg.enable_ce:
......
...@@ -32,10 +32,8 @@ class GTrainer(): ...@@ -32,10 +32,8 @@ class GTrainer():
model = StarGAN_model() model = StarGAN_model()
self.fake_img = model.network_G( self.fake_img = model.network_G(
image_real, label_trg, cfg, name="g_main") image_real, label_trg, cfg, name="g_main")
self.fake_img.persistable = True
self.rec_img = model.network_G( self.rec_img = model.network_G(
self.fake_img, label_org, cfg, name="g_main") self.fake_img, label_org, cfg, name="g_main")
self.rec_img.persistable = True
self.infer_program = self.program.clone(for_test=False) self.infer_program = self.program.clone(for_test=False)
self.g_loss_rec = fluid.layers.reduce_mean( self.g_loss_rec = fluid.layers.reduce_mean(
fluid.layers.abs( fluid.layers.abs(
...@@ -58,9 +56,6 @@ class GTrainer(): ...@@ -58,9 +56,6 @@ class GTrainer():
fluid.layers.sigmoid_cross_entropy_with_logits( fluid.layers.sigmoid_cross_entropy_with_logits(
self.cls_fake, label_trg)) / cfg.batch_size self.cls_fake, label_trg)) / cfg.batch_size
self.g_loss = self.g_loss_fake + cfg.lambda_rec * self.g_loss_rec + self.g_loss_cls self.g_loss = self.g_loss_fake + cfg.lambda_rec * self.g_loss_rec + self.g_loss_cls
self.g_loss_fake.persistable = True
self.g_loss_rec.persistable = True
self.g_loss_cls.persistable = True
lr = cfg.g_lr lr = cfg.g_lr
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
...@@ -94,7 +89,7 @@ class DTrainer(): ...@@ -94,7 +89,7 @@ class DTrainer():
with fluid.program_guard(self.program): with fluid.program_guard(self.program):
model = StarGAN_model() model = StarGAN_model()
image_real = fluid.layers.data( image_real = fluid.data(
name='image_real', shape=image_real.shape, dtype='float32') name='image_real', shape=image_real.shape, dtype='float32')
self.fake_img = model.network_G( self.fake_img = model.network_G(
image_real, label_trg, cfg, name="g_main") image_real, label_trg, cfg, name="g_main")
...@@ -124,10 +119,6 @@ class DTrainer(): ...@@ -124,10 +119,6 @@ class DTrainer():
name="d_main") name="d_main")
self.d_loss = self.d_loss_real + self.d_loss_fake + self.d_loss_cls + cfg.lambda_gp * self.d_loss_gp self.d_loss = self.d_loss_real + self.d_loss_fake + self.d_loss_cls + cfg.lambda_gp * self.d_loss_gp
self.d_loss_real.persistable = True
self.d_loss_fake.persistable = True
self.d_loss_gp.persistable = True
self.d_loss_cls.persistable = True
vars = [] vars = []
for var in self.program.list_vars(): for var in self.program.list_vars():
if fluid.io.is_parameter(var) and var.name.startswith("d_"): if fluid.io.is_parameter(var) and var.name.startswith("d_"):
...@@ -280,7 +271,7 @@ class StarGAN(object): ...@@ -280,7 +271,7 @@ class StarGAN(object):
if self.cfg.enable_ce: if self.cfg.enable_ce:
fluid.default_startup_program().random_seed = 90 fluid.default_startup_program().random_seed = 90
py_reader = fluid.io.PyReader( loader = fluid.io.DataLoader.from_generator(
feed_list=[image_real, label_org, label_trg], feed_list=[image_real, label_org, label_trg],
capacity=128, capacity=128,
iterable=True, iterable=True,
...@@ -293,7 +284,7 @@ class StarGAN(object): ...@@ -293,7 +284,7 @@ class StarGAN(object):
# prepare environment # prepare environment
place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace() place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace()
py_reader.decorate_batch_generator( loader.set_batch_generator(
self.train_reader, self.train_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
...@@ -301,8 +292,8 @@ class StarGAN(object): ...@@ -301,8 +292,8 @@ class StarGAN(object):
exe.run(fluid.default_startup_program()) exe.run(fluid.default_startup_program())
if self.cfg.init_model: if self.cfg.init_model:
utility.init_checkpoints(self.cfg, exe, gen_trainer, "net_G") utility.init_checkpoints(self.cfg, gen_trainer, "net_G")
utility.init_checkpoints(self.cfg, exe, dis_trainer, "net_D") utility.init_checkpoints(self.cfg, dis_trainer, "net_D")
### memory optim ### memory optim
build_strategy = fluid.BuildStrategy() build_strategy = fluid.BuildStrategy()
...@@ -324,7 +315,7 @@ class StarGAN(object): ...@@ -324,7 +315,7 @@ class StarGAN(object):
total_train_batch = 0 # used for benchmark total_train_batch = 0 # used for benchmark
for epoch_id in range(self.cfg.epoch): for epoch_id in range(self.cfg.epoch):
batch_id = 0 batch_id = 0
for data in py_reader(): for data in loader():
if self.cfg.max_iter and total_train_batch == self.cfg.max_iter: # used for benchmark if self.cfg.max_iter and total_train_batch == self.cfg.max_iter: # used for benchmark
return return
s_time = time.time() s_time = time.time()
...@@ -378,24 +369,24 @@ class StarGAN(object): ...@@ -378,24 +369,24 @@ class StarGAN(object):
name='image_name', name='image_name',
shape=[None, self.cfg.n_samples], shape=[None, self.cfg.n_samples],
dtype='int32') dtype='int32')
test_py_reader = fluid.io.PyReader( test_loader = fluid.io.DataLoader.from_generator(
feed_list=[image_real, label_org, label_trg, image_name], feed_list=[image_real, label_org, label_trg, image_name],
capacity=32, capacity=32,
iterable=True, iterable=True,
use_double_buffer=True) use_double_buffer=True)
test_py_reader.decorate_batch_generator( test_loader.set_batch_generator(
self.test_reader, self.test_reader,
places=fluid.cuda_places() places=fluid.cuda_places()
if self.cfg.use_gpu else fluid.cpu_places()) if self.cfg.use_gpu else fluid.cpu_places())
test_program = gen_trainer.infer_program test_program = gen_trainer.infer_program
utility.save_test_image(epoch_id, self.cfg, exe, place, utility.save_test_image(epoch_id, self.cfg, exe, place,
test_program, gen_trainer, test_program, gen_trainer,
test_py_reader) test_loader)
if self.cfg.save_checkpoints: if self.cfg.save_checkpoints:
utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, utility.checkpoints(epoch_id, self.cfg, gen_trainer,
"net_G") "net_G")
utility.checkpoints(epoch_id, self.cfg, exe, dis_trainer, utility.checkpoints(epoch_id, self.cfg, dis_trainer,
"net_D") "net_D")
# used for continuous evaluation # used for continuous evaluation
if self.cfg.enable_ce: if self.cfg.enable_ce:
......
...@@ -48,85 +48,25 @@ def plot(gen_data): ...@@ -48,85 +48,25 @@ def plot(gen_data):
return fig return fig
def checkpoints(epoch, cfg, exe, trainer, name): def checkpoints(epoch, cfg, trainer, name):
output_path = os.path.join(cfg.output, 'checkpoints', str(epoch)) output_path = os.path.join(cfg.output, 'checkpoints', str(epoch))
if not os.path.exists(output_path): if not os.path.exists(output_path):
os.makedirs(output_path) os.makedirs(output_path)
fluid.io.save_persistables( fluid.save(
exe, os.path.join(output_path, name), main_program=trainer.program) trainer.program, os.path.join(output_path, name))
print('save checkpoints {} to {}'.format(name, output_path)) print('save checkpoints {} to {}'.format(name, output_path))
sys.stdout.flush() sys.stdout.flush()
def init_checkpoints(cfg, exe, trainer, name): def init_checkpoints(cfg, trainer, name):
assert os.path.exists(cfg.init_model), "{} cannot be found.".format( assert os.path.exists(cfg.init_model), "{} cannot be found.".format(
cfg.init_model) cfg.init_model)
fluid.io.load_persistables( fluid.load(
exe, os.path.join(cfg.init_model, name), main_program=trainer.program) trainer.program, os.path.join(cfg.init_model, name))
print('load checkpoints {} {} DONE'.format(cfg.init_model, name)) print('load checkpoints {} {} DONE'.format(cfg.init_model, name))
sys.stdout.flush() sys.stdout.flush()
### the initialize checkpoint is one file named checkpoint.pdparams
def init_from_checkpoint(args, exe, trainer, name):
if not os.path.exists(args.init_model):
raise Warning("the checkpoint path does not exist.")
return False
fluid.io.load_persistables(
executor=exe,
dirname=os.path.join(args.init_model, name),
main_program=trainer.program,
filename="checkpoint.pdparams")
print("finish initing model from checkpoint from %s" % (args.init_model))
return True
### save the parameters of generator to one file
def save_param(args, exe, program, dirname, var_name="generator"):
param_dir = os.path.join(args.output, 'infer_vars')
if not os.path.exists(param_dir):
os.makedirs(param_dir)
def _name_has_generator(var):
res = (fluid.io.is_parameter(var) and var.name.startswith(var_name))
print(var.name, res)
return res
fluid.io.save_vars(
exe,
os.path.join(param_dir, dirname),
main_program=program,
predicate=_name_has_generator,
filename="params.pdparams")
print("save parameters at %s" % (os.path.join(param_dir, dirname)))
return True
### save the checkpoint to one file
def save_checkpoint(epoch, args, exe, trainer, dirname):
checkpoint_dir = os.path.join(args.output, 'checkpoints', str(epoch))
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
fluid.io.save_persistables(
exe,
os.path.join(checkpoint_dir, dirname),
main_program=trainer.program,
filename="checkpoint.pdparams")
print("save checkpoint at %s" % (os.path.join(checkpoint_dir, dirname)))
return True
def save_test_image(epoch, def save_test_image(epoch,
cfg, cfg,
exe, exe,
...@@ -384,7 +324,6 @@ def check_attribute_conflict(label_batch, attr, attrs): ...@@ -384,7 +324,6 @@ def check_attribute_conflict(label_batch, attr, attrs):
def save_batch_image(img): def save_batch_image(img):
#if img.shape[0] == 1:
if len(img) == 1: if len(img) == 1:
res_img = np.squeeze(img).transpose([1, 2, 0]) res_img = np.squeeze(img).transpose([1, 2, 0])
else: else:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册