#copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve. # #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.DCGAN_network import DCGAN_model from util import utility import sys import six import os import numpy as np import time import matplotlib matplotlib.use('agg') import matplotlib.pyplot as plt import paddle.fluid as fluid import random class GTrainer(): def __init__(self, input, label, cfg): self.program = fluid.default_main_program().clone() 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( input, dtype='float32', shape=[-1, 1], value=1.0) 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(): if fluid.io.is_parameter(var) and (var.name.startswith("G")): vars.append(var.name) optimizer = fluid.optimizer.Adam( learning_rate=cfg.learning_rate, beta1=0.5, name="net_G") optimizer.minimize(self.g_loss, parameter_list=vars) class DTrainer(): def __init__(self, input, labels, cfg): self.program = fluid.default_main_program().clone() with fluid.program_guard(self.program): model = DCGAN_model(cfg.batch_size) d_logit = model.network_D(input, name="D") 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")): vars.append(var.name) optimizer = fluid.optimizer.Adam( learning_rate=cfg.learning_rate, beta1=0.5, name="net_D") optimizer.minimize(self.d_loss, parameter_list=vars) class DCGAN(object): def add_special_args(self, parser): parser.add_argument( '--noise_size', type=int, default=100, help="the noise dimension") parser.add_argument( '--enable_ce', action='store_true', help="if set, run the tasks with continuous evaluation logs") return parser def __init__(self, cfg=None, train_reader=None): self.cfg = cfg self.train_reader = train_reader def build_model(self): img = fluid.data(name='img', shape=[None, 784], dtype='float32') noise = fluid.data( name='noise', shape=[None, self.cfg.noise_size], dtype='float32') label = fluid.data(name='label', shape=[None, 1], dtype='float32') # used for continuous evaluation if self.cfg.enable_ce: fluid.default_startup_program().random_seed = 90 random.seed(0) np.random.seed(0) g_trainer = GTrainer(noise, label, self.cfg) d_trainer = DTrainer(img, label, self.cfg) # prepare enviorment place = fluid.CUDAPlace(0) if self.cfg.use_gpu else fluid.CPUPlace() exe = fluid.Executor(place) exe.run(fluid.default_startup_program()) const_n = np.random.uniform( low=-1.0, high=1.0, 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") ### memory optim build_strategy = fluid.BuildStrategy() build_strategy.enable_inplace = True g_trainer_program = fluid.CompiledProgram( g_trainer.program).with_data_parallel( loss_name=g_trainer.g_loss.name, build_strategy=build_strategy) d_trainer_program = fluid.CompiledProgram( d_trainer.program).with_data_parallel( loss_name=d_trainer.d_loss.name, build_strategy=build_strategy) if self.cfg.run_test: image_path = os.path.join(self.cfg.output, 'test') if not os.path.exists(image_path): os.makedirs(image_path) t_time = 0 for epoch_id in range(self.cfg.epoch): for batch_id, data in enumerate(self.train_reader()): if len(data) != self.cfg.batch_size: continue noise_data = np.random.uniform( low=-1.0, high=1.0, size=[self.cfg.batch_size, self.cfg.noise_size]).astype( 'float32') real_image = np.array(list(map(lambda x: x[0], data))).reshape( [-1, 784]).astype('float32') real_label = np.ones( shape=[real_image.shape[0], 1], dtype='float32') fake_label = np.zeros( shape=[real_image.shape[0], 1], dtype='float32') s_time = time.time() generate_image = exe.run(g_trainer_program, feed={'noise': noise_data}, fetch_list=[g_trainer.fake]) d_real_loss = exe.run( d_trainer_program, feed={'img': real_image, 'label': real_label}, fetch_list=[d_trainer.d_loss])[0] d_fake_loss = exe.run( d_trainer_program, feed={'img': generate_image[0], 'label': fake_label}, fetch_list=[d_trainer.d_loss])[0] d_loss = d_real_loss + d_fake_loss for _ in six.moves.xrange(self.cfg.num_generator_time): noise_data = np.random.uniform( low=-1.0, high=1.0, size=[self.cfg.batch_size, self.cfg.noise_size]).astype( 'float32') g_loss = exe.run(g_trainer_program, feed={'noise': noise_data}, fetch_list=[g_trainer.g_loss])[0] batch_time = time.time() - s_time if batch_id % self.cfg.print_freq == 0: print( 'Epoch ID: {} Batch ID: {} D_loss: {} G_loss: {} Batch_time_cost: {}'. format(epoch_id, batch_id, d_loss[0], g_loss[0], batch_time)) t_time += batch_time if self.cfg.run_test: generate_const_image = exe.run( g_trainer.infer_program, feed={'noise': const_n}, fetch_list=[g_trainer.fake])[0] generate_image_reshape = np.reshape(generate_const_image, ( self.cfg.batch_size, -1)) total_images = np.concatenate( [real_image, generate_image_reshape]) fig = utility.plot(total_images) plt.title('Epoch ID={}, Batch ID={}'.format(epoch_id, batch_id)) img_name = '{:04d}_{:04d}.png'.format(epoch_id, batch_id) plt.savefig( os.path.join(image_path, img_name), bbox_inches='tight') 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") # used for continuous evaluation if self.cfg.enable_ce: device_num = fluid.core.get_cuda_device_count( ) if self.cfg.use_gpu else 1 print("kpis\tdcgan_d_loss_card{}\t{}".format(device_num, d_loss[0])) print("kpis\tdcgan_g_loss_card{}\t{}".format(device_num, g_loss[0])) print("kpis\tdcgan_Batch_time_cost_card{}\t{}".format(device_num, batch_time))