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

Sync gan 1.7 (#4288)

* sync_1.7

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