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
