# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import from __future__ import division from __future__ import print_function from network.StarGAN_network import StarGAN_model from util import utility import paddle.fluid as fluid import sys import time import copy import numpy as np import pickle as pkl class GTrainer(): def __init__(self, image_real, label_org, label_trg, cfg, step_per_epoch): self.program = fluid.default_main_program().clone() with fluid.program_guard(self.program): 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( fluid.layers.elementwise_sub( x=image_real, y=self.rec_img))) self.pred_fake, self.cls_fake = model.network_D( self.fake_img, cfg, name="d_main") if cfg.gan_mode != 'wgan': raise NotImplementedError( "gan_mode {} is not support! only support wgan".format( cfg.gan_mode)) #wgan self.g_loss_fake = -1 * fluid.layers.mean(self.pred_fake) cls_shape = self.cls_fake.shape self.cls_fake = fluid.layers.reshape( self.cls_fake, [-1, cls_shape[1] * cls_shape[2] * cls_shape[3]]) self.g_loss_cls = fluid.layers.reduce_sum( 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(): if fluid.io.is_parameter(var) and var.name.startswith("g_"): vars.append(var.name) self.param = vars total_iters = step_per_epoch * cfg.epoch boundaries = [cfg.num_iters - cfg.num_iters_decay] values = [lr] for x in range(cfg.num_iters - cfg.num_iters_decay + 1, total_iters): if x % cfg.lr_update_step == 0: boundaries.append(x) lr -= (lr / float(cfg.num_iters_decay)) values.append(lr) lr = values[-1] lr -= (lr / float(cfg.num_iters_decay)) values.append(lr) optimizer = fluid.optimizer.Adam( learning_rate=fluid.layers.piecewise_decay( boundaries=boundaries, values=values), beta1=0.5, beta2=0.999, name="net_G") optimizer.minimize(self.g_loss, parameter_list=vars) class DTrainer(): def __init__(self, image_real, label_org, label_trg, cfg, step_per_epoch): self.program = fluid.default_main_program().clone() with fluid.program_guard(self.program): model = StarGAN_model() image_real = fluid.layers.data( name='image_real', shape=image_real.shape, dtype='float32') self.fake_img = model.network_G( image_real, label_trg, cfg, name="g_main") self.pred_real, self.cls_real = model.network_D( image_real, cfg, name="d_main") self.pred_fake, _ = model.network_D( self.fake_img, cfg, name="d_main") cls_shape = self.cls_real.shape self.cls_real = fluid.layers.reshape( self.cls_real, [-1, cls_shape[1] * cls_shape[2] * cls_shape[3]]) self.d_loss_cls = fluid.layers.reduce_sum( fluid.layers.sigmoid_cross_entropy_with_logits( self.cls_real, label_org)) / cfg.batch_size if cfg.gan_mode != 'wgan': raise NotImplementedError( "gan_mode {} is not support! only support wgan".format( cfg.gan_mode)) #wgan self.d_loss_fake = fluid.layers.mean(self.pred_fake) self.d_loss_real = -1 * fluid.layers.mean(self.pred_real) self.d_loss_gp = self.gradient_penalty( getattr(model, "network_D"), image_real, self.fake_img, cfg=cfg, 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_"): vars.append(var.name) self.param = vars total_iters = step_per_epoch * cfg.epoch boundaries = [cfg.num_iters - cfg.num_iters_decay] values = [cfg.d_lr] lr = cfg.d_lr for x in range(cfg.num_iters - cfg.num_iters_decay + 1, total_iters): if x % cfg.lr_update_step == 0: boundaries.append(x) lr -= (lr / float(cfg.num_iters_decay)) values.append(lr) lr = values[-1] lr -= (lr / float(cfg.num_iters_decay)) values.append(lr) optimizer = fluid.optimizer.Adam( learning_rate=fluid.layers.piecewise_decay( boundaries=boundaries, values=values), beta1=0.5, beta2=0.999, name="net_D") optimizer.minimize(self.d_loss, parameter_list=vars) def gradient_penalty(self, f, real, fake, cfg=None, name=None): def _interpolate(a, b): shape = [a.shape[0]] alpha = fluid.layers.uniform_random_batch_size_like( input=a, shape=shape, min=0.0, max=1.0) inner = b * (1.0 - alpha) + a * alpha return inner x = _interpolate(real, fake) pred, _ = f(x, cfg, name=name) if isinstance(pred, tuple): pred = pred[0] vars = [] for var in fluid.default_main_program().list_vars(): if fluid.io.is_parameter(var) and var.name.startswith('d_'): vars.append(var.name) grad = fluid.gradients(pred, x, no_grad_set=vars)[0] grad_shape = grad.shape grad = fluid.layers.reshape( grad, [-1, grad_shape[1] * grad_shape[2] * grad_shape[3]]) epsilon = 1e-16 norm = fluid.layers.sqrt( fluid.layers.reduce_sum( fluid.layers.square(grad), dim=1) + epsilon) gp = fluid.layers.reduce_mean(fluid.layers.square(norm - 1.0)) return gp class StarGAN(object): def add_special_args(self, parser): parser.add_argument( '--g_lr', type=float, default=0.0001, help="learning rate of g") parser.add_argument( '--d_lr', type=float, default=0.0001, help="learning rate of d") parser.add_argument( '--c_dim', type=int, default=5, help="the number of attributes we selected") parser.add_argument( '--g_repeat_num', type=int, default=6, help="number of layers in generator") parser.add_argument( '--d_repeat_num', type=int, default=6, help="number of layers in discriminator") parser.add_argument( '--num_iters', type=int, default=200000, help="num iters") parser.add_argument( '--num_iters_decay', type=int, default=100000, help="num iters decay") parser.add_argument( '--lr_update_step', type=int, default=1000, help="iters when lr update ") parser.add_argument( '--lambda_cls', type=float, default=1.0, help="the coefficient of classification") parser.add_argument( '--lambda_rec', type=float, default=10.0, help="the coefficient of refactor") parser.add_argument( '--lambda_gp', type=float, default=10.0, help="the coefficient of gradient penalty") parser.add_argument( '--n_critic', type=int, default=5, help="discriminator training steps when generator update") parser.add_argument( '--selected_attrs', type=str, default="Black_Hair,Blond_Hair,Brown_Hair,Male,Young", help="the attributes we selected to change") parser.add_argument( '--n_samples', type=int, default=1, help="batch size when testing") return parser def __init__(self, cfg=None, train_reader=None, test_reader=None, batch_num=1, id2name=None): self.cfg = cfg self.train_reader = train_reader self.test_reader = test_reader self.batch_num = batch_num def build_model(self): data_shape = [-1, 3, self.cfg.image_size, self.cfg.image_size] image_real = fluid.layers.data( name='image_real', shape=data_shape, dtype='float32') label_org = fluid.layers.data( name='label_org', shape=[self.cfg.c_dim], dtype='float32') label_trg = fluid.layers.data( name='label_trg', shape=[self.cfg.c_dim], dtype='float32') py_reader = fluid.io.PyReader( feed_list=[image_real, label_org, label_trg], capacity=128, iterable=True, use_double_buffer=True) gen_trainer = GTrainer(image_real, label_org, label_trg, self.cfg, self.batch_num) dis_trainer = DTrainer(image_real, label_org, label_trg, self.cfg, self.batch_num) # prepare environment place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace() py_reader.decorate_batch_generator( self.train_reader, places=fluid.cuda_places() if self.cfg.use_gpu else fluid.cpu_places()) exe = fluid.Executor(place) 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") ### memory optim build_strategy = fluid.BuildStrategy() gen_trainer_program = fluid.CompiledProgram( gen_trainer.program).with_data_parallel( loss_name=gen_trainer.g_loss.name, build_strategy=build_strategy) dis_trainer_program = fluid.CompiledProgram( dis_trainer.program).with_data_parallel( loss_name=dis_trainer.d_loss.name, build_strategy=build_strategy) t_time = 0 for epoch_id in range(self.cfg.epoch): batch_id = 0 for data in py_reader(): s_time = time.time() d_loss_real, d_loss_fake, d_loss, d_loss_cls, d_loss_gp = exe.run( dis_trainer_program, fetch_list=[ dis_trainer.d_loss_real, dis_trainer.d_loss_fake, dis_trainer.d_loss, dis_trainer.d_loss_cls, dis_trainer.d_loss_gp ], feed=data) # optimize the generator network if (batch_id + 1) % self.cfg.n_critic == 0: g_loss_fake, g_loss_rec, g_loss_cls, fake_img, rec_img = exe.run( gen_trainer_program, fetch_list=[ gen_trainer.g_loss_fake, gen_trainer.g_loss_rec, gen_trainer.g_loss_cls, gen_trainer.fake_img, gen_trainer.rec_img ], feed=data) print("epoch{}: batch{}: \n\ g_loss_fake: {}; g_loss_rec: {}; g_loss_cls: {}" .format(epoch_id, batch_id, g_loss_fake[0], g_loss_rec[0], g_loss_cls[0])) batch_time = time.time() - s_time t_time += batch_time if (batch_id + 1) % self.cfg.print_freq == 0: print("epoch{}: batch{}: \n\ d_loss_real: {}; d_loss_fake: {}; d_loss_cls: {}; d_loss_gp: {} \n\ Batch_time_cost: {}".format( epoch_id, batch_id, d_loss_real[0], d_loss_fake[ 0], d_loss_cls[0], d_loss_gp[0], batch_time)) sys.stdout.flush() batch_id += 1 if self.cfg.run_test: image_name = fluid.layers.data( name='image_name', shape=[self.cfg.n_samples], dtype='int32') test_py_reader = fluid.io.PyReader( feed_list=[image_real, label_org, label_trg, image_name], capacity=32, iterable=True, use_double_buffer=True) test_py_reader.decorate_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) if self.cfg.save_checkpoints: utility.checkpoints(epoch_id, self.cfg, exe, gen_trainer, "net_G") utility.checkpoints(epoch_id, self.cfg, exe, dis_trainer, "net_D")