def example_gan(adversarial_optimizer, path, opt_g, opt_d, nb_epoch, generator, discriminator, latent_dim, targets=gan_targets, loss="binary_crossentropy"): csvpath = os.path.join(path, "history.csv") if os.path.exists(csvpath): print("Already exists: {}".format(csvpath)) return print("Training: {}".format(csvpath)) # gan (x - > yfake, yreal), z is gaussian generated on GPU # can also experiment with uniform_latent_sampling d_g = discriminator(0) d_d = discriminator(0.5) generator.summary() d_d.summary() gan_g = simple_gan(generator, d_g, None) gan_d = simple_gan(generator, d_d, None) x = gan_g.inputs[1] z = normal_latent_sampling((latent_dim, ))(x) # estiminate z from inputs gan_g = Model([x], fix_names(gan_g([z, x]), gan_g.output_names)) gan_d = Model([x], fix_names(gan_d([z, x]), gan_d.output_names)) # build adversarial model model = AdversarialModel( player_models=[gan_g, gan_d], player_params=[generator.trainable_weights, d_d.trainable_weights], player_names=["generator", "discriminator"]) model.adversarial_compile(adversarial_optimizer=adversarial_optimizer, player_optimizers=[opt_g, opt_d], loss=loss) # create callback to generate images zsamples = np.random.normal(size=(10 * 10, latent_dim)) def generator_sampler(): xpred = generator.predict(zsamples) xpred = dim_ordering_unfix(xpred.transpose((0, 2, 3, 1))) return xpred.reshape((10, 10) + xpred.shape[1:]) generator_cb = ImageGridCallback(os.path.join(path, "epoch-{:03d}.png"), generator_sampler, cmap=None) callbacks = [generator_cb] if K.backend() == "tensorflow": callbacks.append( TensorBoard(log_dir=os.path.join(path, "logs"), histogram_freq=0, write_graph=True, write_images=True)) # train model """ xtrain, xtest = cifar10_data() model.add(UpSampling2D(size=(2, 2))) y = targets(xtrain.shape[0]) ytest = targets(xtest.shape[0]) history = model.fit(x=xtrain, y=y, validation_data=(xtest, ytest), callbacks=callbacks, epochs=nb_epoch, batch_size=32) """ # train_generator, validation_generator = ioswp_data() train_generator, validation_generator, xtrain, xtest = ioswp_data() y = targets(xtrain.shape[0]) ytest = targets(xtest.shape[0]) history = model.fit_generator( train_generator, # validation_data=(xtest, ytest), validation_data=validation_generator, callbacks=callbacks, epochs=nb_epoch, steps_per_epoch=(143 / 13) * 1000, validation_steps=(143 / 13) * 100) # save history to CSV df = pd.DataFrame(history.history) df.to_csv(csvpath) # save models generator.save(os.path.join(path, "generator.h5")) d_d.save(os.path.join(path, "discriminator.h5"))
def example_faae(path, adversarial_optimizer): latent_dim = 256 units = 512 input_shape = dim_ordering_shape((3, 32, 32)) # generator (z -> x) generator = model_generator(latent_dim, units=units) # encoder (x ->z) encoder = model_encoder(latent_dim, input_shape, units=units) # autoencoder (x -> x') autoencoder = Model(encoder.inputs, generator(encoder(encoder.inputs))) # discriminator (z -> y) discriminator = model_discriminator() # build FAAE zreal = discriminator.inputs[0] x = generator.inputs[0] z = generator(x) xpred = encoder(z) yreal = discriminator(zreal) yfake = discriminator(z) aae = Model([zreal, x], fix_names([xpred, yfake, yreal], ["xpred", "yfake", "yreal"])) # print summary of models generator.summary() encoder.summary() discriminator.summary() #encoder.load_weights(os.path.join(path, "encoder.h5")) #generator.load_weights(os.path.join(path, "generator.h5")) #discriminator.load_weights(os.path.join(path, "discriminator.h5")) # build adversarial model generative_params = generator.trainable_weights + encoder.trainable_weights model = AdversarialModel( base_model=aae, player_params=[generative_params, discriminator.trainable_weights], player_names=["generator", "discriminator"]) model.adversarial_compile( adversarial_optimizer=adversarial_optimizer, player_optimizers=[Adam(3e-4, decay=1e-4), Adam(1e-3, decay=1e-4)], loss={ "yfake": "binary_crossentropy", "yreal": "binary_crossentropy", "xpred": "mean_squared_error" }, player_compile_kwargs=[{ "loss_weights": { "yfake": 1, "yreal": 1, "xpred": 8 } }] * 2) xtrain, xtest = cifar10_data() def generator_sampler(): zsamples = np.random.randn(10 * 10, latent_dim) return dim_ordering_unfix(generator.predict(zsamples)).transpose( (0, 2, 3, 1)).reshape((10, 10, 32, 32, 3)) generator_cb = ImageGridCallback( os.path.join(path, "generated-epoch-{:03d}.png"), generator_sampler) def autoencoder_sampler(): xsamples = n_choice(xtest, 10) xrep = np.repeat(xsamples, 9, axis=0) xgen = dim_ordering_unfix(autoencoder.predict(xrep)).reshape( (10, 9, 3, 32, 32)) xsamples = dim_ordering_unfix(xsamples).reshape((10, 1, 3, 32, 32)) samples = np.concatenate((xsamples, xgen), axis=1) samples = samples.transpose((0, 1, 3, 4, 2)) return samples autoencoder_cb = ImageGridCallback(os.path.join( path, "autoencoded-epoch-{:03d}.png"), autoencoder_sampler, cmap=None) train_datagen = gen_sample(128, 256, False) test_datagen = gen_sample(32, 256, True) history = model.fit_generator(train_datagen, epochs=200, steps_per_epoch=1000, validation_data=test_datagen, validation_steps=100, callbacks=[generator_cb, autoencoder_cb]) # save history df = pd.DataFrame(history.history) df.to_csv(os.path.join(path, "history.csv")) # save model encoder.save(os.path.join(path, "encoder.h5")) generator.save(os.path.join(path, "generator.h5")) discriminator.save(os.path.join(path, "discriminator.h5"))
class UnifAI(object): def __init__(self, config): self.config = config self.model_config = self.config.model_config # alias self.module_builder = UnifAI_ModuleBuilder( self.model_config, self.config.remote_weights_path ) self.encoder = None self.predictor = None self.decoder = None self.noisy_transformer = None self.disentangler1 = None self.disentangler2 = None self.z_discriminator = None self.model_inference = None self.model_train = None self.optimizers = self.config.optimizers # alias self.compiled = False def __random_target(self, x, dim, embedding_activation): range_low = -1 if embedding_activation == 'tanh' else 0 range_high = 1 tfake = uniform_latent_sampling( (dim,), low=range_low, high=range_high )(x) return tfake def __build_connected_network_train(self, main=True): x = self.encoder.inputs[0] e1, e2 = self.encoder(x) noisy_e1 = self.noisy_transformer(e1) y = self.predictor(e1) x_pred = self.decoder([noisy_e1, e2]) e1_dim = int(self.encoder.outputs[0].shape[-1]) e2_dim = int(self.encoder.outputs[1].shape[-1]) output_vars = [y, x_pred] output_names = ['y', 'x_pred'] e1_target = e1 e2_target = e2 e2_pred = self.disentangler1(e1) e1_pred = self.disentangler2(e2) if main: embedding_activation = self.model_config.embedding_activation e2_target = self.__random_target(x, e2_dim, embedding_activation) e1_target = self.__random_target(x, e1_dim, embedding_activation) e1_e1_pred = Concatenate()([e1_target, e1_pred]) output_vars.append(e1_e1_pred) output_names.append('e1pred') e2_e2_pred = Concatenate()([e2_target, e2_pred]) output_vars.append(e2_e2_pred) output_names.append('e2pred') if self.z_discriminator is not None: z = self.z_discriminator(e1) output_vars.append(z) output_names.append('z') outputs = fix_names(output_vars, output_names) network = Model(inputs=[x], outputs=outputs) return network def build_model_train(self): if self.model_train is None: with tf.device('/gpu:0'): # Build modules: ## Encoder: x -> [e1, e2] ## Predictor: e1 -> y ## Noisy-transformer: e1 -> noisy_e1 ## Decoder: [noisy_e1, e2] -> x self.encoder, self.predictor, self.decoder = \ self.module_builder.build_default_modules( ['encoder', 'predictor', 'decoder'] ) self.noisy_transformer = self.module_builder.build_module( 'noisy_transformer', name='noisy_transformer', build_kwargs={ 'params': [self.config.dropout_rate] } ) ## Disentanglers: self.disentangler1 = self.module_builder.build_module( 'disentangler', name='disentangler1', build_kwargs={ 'input_dim': self.model_config.embedding_dim_1, 'output_dim': self.model_config.embedding_dim_2 } ) self.disentangler2 = self.module_builder.build_module( 'disentangler', name='disentangler2', build_kwargs={ 'input_dim': self.model_config.embedding_dim_2, 'output_dim': self.model_config.embedding_dim_1 } ) ## z_discriminator: if self.config.bias: self.z_discriminator = self.module_builder.build_module( 'z_discriminator', name='z_discriminator' ) # Build 2 copies of the connected network main_model = self.__build_connected_network_train(main=True) adv_model = self.__build_connected_network_train(main=False) models = [main_model, adv_model] # Parallelize over GPUs if self.config.num_gpus > 1: for i in range(len(models)): models[i] = \ multi_gpu_model(models[i], gpus=self.config.num_gpus) # Create final model ## Gather params main_params = self.encoder.trainable_weights + \ self.predictor.trainable_weights + self.decoder.trainable_weights adv_params = self.disentangler1.trainable_weights + \ self.disentangler2.trainable_weights if self.config.bias: adv_params.extend(self.z_discriminator.trainable_weights) ## Build keras_adversarial model self.model_train = AdversarialModel( player_models=models, player_params=[main_params, adv_params], player_names=['main_model', 'adv_model'] ) def compile_model(self): assert self.model_train is not None, 'run build_model_train()' optimizers = self.config.optimizers losses = self.config.losses main_loss_weights = [lw for lw in self.config.loss_weights] adv_loss_weights = [lw for lw in self.config.loss_weights] player_compile_kwargs = [ { 'loss_weights': main_loss_weights, 'metrics': self.config.metrics }, { 'loss_weights': adv_loss_weights, 'metrics': self.config.metrics } ] adversarial_optimizer = AdversarialOptimizerScheduled( [int(p) for p in list(self.config.training_schedule)] ) self.model_train.adversarial_compile( adversarial_optimizer=adversarial_optimizer, player_optimizers=optimizers, loss=losses, player_compile_kwargs=player_compile_kwargs ) self.compiled = True def build_compiled_model(self): self.build_model_train() self.compile_model() def fit(self, dtrain, dvalid, streaming_data=False, epochs=None, callbacks=None, training_steps=None, validation_steps=None): assert self.compiled, 'run compile_model() before training' if streaming_data: train_generator = dtrain valid_generator = dvalid self.model_train.fit_generator( train_generator, steps_per_epoch=training_steps, validation_data=valid_generator, validation_steps=validation_steps, callbacks=callbacks, epochs=epochs ) else: xtrain, ytrain = dtrain xvalid, yvalid = dvalid self.model_train.fit( x=xtrain, y=ytrain, validation_data=(xvalid, yvalid), callbacks=callbacks, epochs=epochs, batch_size=(self.config.batch_size * self.config.num_gpus) ) def build_model_inference(self, checkpoint_epoch=None): if self.model_inference is None: device = '/cpu:0' if self.config.num_gpus > 1 else '/gpu:0' with tf.device(device): self.encoder = self.module_builder.build_module( 'encoder', epoch=checkpoint_epoch ) self.predictor = self.module_builder.build_module( 'predictor', epoch=checkpoint_epoch ) x = self.encoder.inputs[0] e1, _ = self.encoder(x) y = self.predictor(e1) self.model_inference = Model(x, y) # Parallelize over GPUs if self.config.num_gpus > 1: self.model_inference = multi_gpu_model( self.model_inference, gpus=self.config.num_gpus ) def predict(self, data, streaming_data=False, prediction_steps=None): assert self.model_inference is not None, 'run build_model_inference() first' if streaming_data: return self.model_inference.predict_generator(data, steps=prediction_steps) else: return self.model_inference.predict(data)
def train_em_gan(adversarial_optimizer, generator, discriminator, gen_opt, disc_opt, latent_dim, h5_filename, h5_dataset_path, sample_shape, output_directory, verbose=1, loss='mean_squared_error', epochs=10, per_epoch=100, r_id="em-gan", is_large_model=False): gan = simple_gan(generator, discriminator, normal_latent_sampling((latent_dim, ))) if verbose >= 1: util.print_model_summaries(generator, discriminator, gan) model = AdversarialModel(base_model=gan, player_params=[ generator.trainable_weights, discriminator.trainable_weights ], player_names=["generator", "discriminator"]) model.adversarial_compile(adversarial_optimizer=adversarial_optimizer, player_optimizers=[gen_opt, disc_opt], loss=loss) zsamples = np.random.normal(size=(5, latent_dim)) sample_generator = util.h5_block_generator(h5_filename, h5_dataset_path, sample_shape, [1, 0, 0, 1]) def generator_sampler(): return generator.predict(zsamples) sampler = util.SampleEM(output_directory, generator_sampler, is_large_model) gen_saver = util.SaveModel(generator, os.path.join(output_directory, "generator")) disc_saver = util.SaveModel( discriminator, os.path.join(output_directory, "discriminator")) history = model.fit_generator(sample_generator, per_epoch, epochs=epochs, verbose=verbose, callbacks=[sampler, gen_saver, disc_saver], validation_data=sample_generator, validation_steps=(per_epoch // 5)) df = pd.DataFrame(history.history) df.to_csv(os.path.join(output_directory, "history.csv")) discriminator.save( os.path.join( output_directory, "gan_disc_" + str(epochs) + "_" + str(per_epoch) + "_" + r_id + ".h5")) generator.save( os.path.join( output_directory, "gan_gen_" + str(epochs) + "_" + str(per_epoch) + "_" + r_id + ".h5")) del model del discriminator del generator