def get_optimizer(optimizer_name="adam", **kwargs):
"""Instantiate the optimizer.
Parameters
----------
optimizer_name : str
Name of the optimizer to use.
Returns
-------
optimizer : tf.keras.optimizers
Optimizer instance used in the model.
"""
# TODO use tensorflow optimizer
if optimizer_name == "adam":
optimizer = Adam(**kwargs)
elif optimizer_name == "adadelta":
optimizer = Adadelta(**kwargs)
elif optimizer_name == "adagrad":
optimizer = Adagrad(**kwargs)
elif optimizer_name == "adamax":
optimizer = Adamax(**kwargs)
elif optimizer_name == "sgd":
optimizer = SGD(**kwargs)
else:
raise ValueError(
"Instead of {0}, optimizer must be chosen among "
"['adam', 'adadelta', 'adagrad', adamax', sgd'].".format(
optimizer_name))
return optimizer
def train(self):
if self.compile_model:
optimizer = Adamax(lr=0.008)
self.model.compile(loss=self.metric,
optimizer=optimizer,
metrics=self.metrics)
# create ordered queues
train_enqueuer = OrderedEnqueuer(self.training_sequence,
use_multiprocessing=False,
shuffle=True)
train_enqueuer.start(workers=1, max_queue_size=64)
train_gen = train_enqueuer.get()
test_enqueuer = OrderedEnqueuer(self.validation_sequence,
use_multiprocessing=False,
shuffle=True)
test_enqueuer.start(workers=1, max_queue_size=64)
test_gen = test_enqueuer.get()
# train model
history = self.model.fit_generator(
generator=train_gen,
epochs=self.epochs,
validation_data=next(test_gen),
verbose=2,
steps_per_epoch=len(self.training_sequence),
validation_steps=len(self.validation_sequence),
callbacks=self.callbacks,
initial_epoch=self.initial_epoch)
# plot metrics
self.plot_results(history, self.output_dir)
self.model.save_weights(str(Path(self.output_dir, 'final_weights')))
def compile(self,
loss: str = 'mean_squared_error',
lr: float = 1e-2,
**params) -> None:
opt = Adamax(learning_rate=lr)
compile_params = {'optimizer': opt, 'metrics': self.metrics, **params}
self.model.compile(loss=loss, **compile_params)
return None
def get_optimizer(self):
if self.optimizer_type == 'adam':
opt = Adam(lr=self.learning_rate)
elif self.optimizer_type == 'adamax':
opt = Adamax(lr=self.learning_rate)
elif self.optimizer_type == 'radam':
opt = RAdam(lr=0.01, total_steps=5000, warmup_proportion=0.1, min_lr=0.0001)
else:
raise ValueError('Undefined OPTIMIZER_TYPE!')
return opt
def get_auto_encoder(self):
model_params = DotDict(get_default_hparams())
auto_encoder = NVAEAutoEncoder64(
model_params,
batch_size=1,
encoder_frames_no=self.encoder_frames_no)
auto_encoder.summary()
metrics = [
auto_encoder.loss_func, auto_encoder.face_metric,
auto_encoder.face_kl_loss, auto_encoder.mask_mse_loss, "mae", "mse"
]
optimizer = Adamax(auto_encoder.hps.learning_rate)
auto_encoder.model.compile(loss=auto_encoder.loss_func,
optimizer=optimizer,
metrics=metrics)
if self.checkpoint_path:
# Load weights:
load_status = auto_encoder.model.load_weights(self.checkpoint_path)
if load_status:
load_status.assert_consumed()
return auto_encoder
def train_big_vae(train_directory,
test_directory,
samples_directory,
epochs=100,
batch_size=32,
encoder_frames_no=45,
initial_epoch=0,
checkpoint_path=None):
encoder_frames_no = encoder_frames_no
input_shape = (128, 128, 3)
train_seq = NVAESequence(train_directory,
input_size=input_shape[:-1],
batch_size=batch_size,
encoder_frames_no=encoder_frames_no)
test_seq = NVAESequence(test_directory,
input_size=input_shape[:-1],
batch_size=batch_size,
encoder_frames_no=encoder_frames_no)
num_samples = 3 # samples to be generated each epoch
model_params = DotDict(get_default_hparams())
model_params["mask_kl_weight"] = 0.1
auto_encoder = BigNVAEAutoEncoder128(model_params,
batch_size=batch_size,
encoder_frames_no=None)
auto_encoder.summary()
model = auto_encoder.model
metric = auto_encoder.loss_func
metrics = [
auto_encoder.loss_func, auto_encoder.face_metric,
auto_encoder.face_kl_loss, auto_encoder.mask_mse_loss, "mae", "mse"
]
optimizer = Adamax(auto_encoder.hps.learning_rate)
model.compile(loss=metric, optimizer=optimizer, metrics=metrics)
callbacks_dict = get_callbacks_dict(auto_encoder, model_params, test_seq,
batch_size, num_samples,
samples_directory)
if checkpoint_path is not None:
# Load weights:
load_status = model.load_weights(checkpoint_path)
load_status.assert_consumed()
# Initial batch (affects LR and KL weight decay):
num_batches = len(train_seq)
count = initial_epoch * num_batches
callbacks_dict['lr_schedule'].count = count
callbacks_dict['kl_weight_schedule'].count = count
callbacks_dict['mask_kl_weight_schedule'].count = count
callbacks = [callback for callback in callbacks_dict.values()]
t = Training(model=model,
training_sequence=train_seq,
validation_sequence=test_seq,
metric=metric,
metrics=metrics,
callbacks=callbacks,
output_dir=samples_directory,
epochs=epochs,
compile_model=False,
initial_epoch=initial_epoch)
t.train()