__author__ = "Lukas Leufen, Felix Kleinert"
__date__ = '2019-12-05'
import json
import logging
import os
import pickle
import keras
from src.data_handling.data_distributor import Distributor
from src.model_modules.keras_extensions import LearningRateDecay, ModelCheckpointAdvanced, CallbackHandler
from src.plotting.training_monitoring import PlotModelHistory, PlotModelLearningRate
from src.run_modules.run_environment import RunEnvironment
class Training(RunEnvironment):
def __init__(self):
super().__init__()
self.model: keras.Model = self.data_store.get("model", "general.model")
self.train_set = None
self.val_set = None
self.test_set = None
self.batch_size = self.data_store.get("batch_size", "general.model")
self.epochs = self.data_store.get("epochs", "general.model")
self.callbacks: CallbackHandler = self.data_store.get("callbacks", "general.model")
self.experiment_name = self.data_store.get("experiment_name", "general")
self._trainable = self.data_store.get("trainable", "general")
self._create_new_model = self.data_store.get("create_new_model", "general")
self._run()
def _run(self) -> None:
"""
Perform training
1) set_generators():
set generators for training, validation and testing and distribute according to batch size
2) make_predict_function():
create predict function before distribution on multiple nodes (detailed information in method description)
3) train():
start or resume training of model and save callbacks
4) save_model():
save best model from training as final model
"""
self.set_generators()
self.make_predict_function()
if self._trainable:
self.train()
self.save_model()
else:
logging.info("No training has started, because trainable parameter was false.")
def make_predict_function(self) -> None:
"""
Creates the predict function. Must be called before distributing. This is necessary, because tf will compile
the predict function just in the moment it is used the first time. This can cause problems, if the model is
distributed on different workers. To prevent this, the function is pre-compiled. See discussion @
https://stackoverflow.com/questions/40850089/is-keras-thread-safe/43393252#43393252
"""
self.model._make_predict_function()
def _set_gen(self, mode: str) -> None:
"""
Set and distribute the generators for given mode regarding batch size
:param mode: name of set, should be from ["train", "val", "test"]
"""
gen = self.data_store.get("generator", f"general.{mode}")
permute_data = self.data_store.get_default("permute_data", f"general.{mode}", default=False)
setattr(self, f"{mode}_set", Distributor(gen, self.model, self.batch_size, permute_data=permute_data))
def set_generators(self) -> None:
"""
Set all generators for training, validation, and testing subsets. The called sub-method will automatically
distribute the data according to the batch size. The subsets can be accessed as class variables train_set,
val_set, and test_set .
"""
for mode in ["train", "val", "test"]:
self._set_gen(mode)
def train(self) -> None:
"""
Perform training using keras fit_generator(). Callbacks are stored locally in the experiment directory. Best
model from training is saved for class variable model. If the file path of checkpoint is not empty, this method
assumes, that this is not a new training starting from the very beginning, but a resumption from a previous
started but interrupted training (or a stopped and now continued training). Train will automatically load the
locally stored information and the corresponding model and proceed with the already started training.
"""
logging.info(f"Train with {len(self.train_set)} mini batches.")
checkpoint = self.callbacks.get_checkpoint()
if not os.path.exists(checkpoint.filepath) or self._create_new_model:
history = self.model.fit_generator(generator=self.train_set.distribute_on_batches(),
steps_per_epoch=len(self.train_set),
epochs=self.epochs,
verbose=2,
validation_data=self.val_set.distribute_on_batches(),
validation_steps=len(self.val_set),
callbacks=self.callbacks.get_callbacks(as_dict=False))
else:
logging.info("Found locally stored model and checkpoints. Training is resumed from the last checkpoint.")
self.callbacks.load_callbacks()
self.callbacks.update_checkpoint()
self.model = keras.models.load_model(checkpoint.filepath)
hist = self.callbacks.get_callback_by_name("hist")
initial_epoch = max(hist.epoch) + 1
_ = self.model.fit_generator(generator=self.train_set.distribute_on_batches(),
steps_per_epoch=len(self.train_set),
epochs=self.epochs,
verbose=2,
validation_data=self.val_set.distribute_on_batches(),
validation_steps=len(self.val_set),
callbacks=self.callbacks.get_callbacks(as_dict=False),
initial_epoch=initial_epoch)
history = hist
try:
lr = self.callbacks.get_callback_by_name("lr")
except IndexError:
lr = None
self.save_callbacks_as_json(history, lr)
self.load_best_model(checkpoint.filepath)
self.create_monitoring_plots(history, lr)
def save_model(self) -> None:
"""
save model in local experiment directory. Model is named as <experiment_name>_<custom_model_name>.h5 .
"""
model_name = self.data_store.get("model_name", "general.model")
logging.debug(f"save best model to {model_name}")
self.model.save(model_name)
self.data_store.set("best_model", self.model, "general")
def load_best_model(self, name: str) -> None:
"""
Load model weights for model with name. Skip if no weights are available.
:param name: name of the model to load weights for
"""
logging.debug(f"load best model: {name}")
try:
self.model.load_weights(name)
logging.info('reload weights...')
except OSError:
logging.info('no weights to reload...')
def save_callbacks_as_json(self, history: keras.callbacks.History, lr_sc: keras.callbacks) -> None:
"""
Save callbacks (history, learning rate) of training.
* history.history -> history.json
* lr_sc.lr -> history_lr.json
:param history: history object of training
"""
logging.debug("saving callbacks")
path = self.data_store.get("experiment_path", "general")
with open(os.path.join(path, "history.json"), "w") as f:
json.dump(history.history, f)
if lr_sc:
with open(os.path.join(path, "history_lr.json"), "w") as f:
json.dump(lr_sc.lr, f)
def create_monitoring_plots(self, history: keras.callbacks.History, lr_sc: LearningRateDecay) -> None:
"""
Creates the history and learning rate plot in dependence of the number of epochs. The plots are saved in the
experiment's plot_path. History plot is named '<exp_name>_history_loss_val_loss.pdf', the learning rate with
'<exp_name>_history_learning_rate.pdf'.
:param history: keras history object with losses to plot (must include 'loss' and 'val_loss')
:param lr_sc: learning rate decay object with 'lr' attribute
"""
path = self.data_store.get("plot_path", "general")
name = self.data_store.get("experiment_name", "general")
# plot history of loss and mse (if available)
filename = os.path.join(path, f"{name}_history_loss.pdf")
PlotModelHistory(filename=filename, history=history)
multiple_branches_used = len(history.model.output_names) > 1 # means that there are multiple output branches
if multiple_branches_used:
filename = os.path.join(path, f"{name}_history_main_loss.pdf")
PlotModelHistory(filename=filename, history=history, main_branch=True)
if "mean_squared_error" in history.model.metrics_names:
filename = os.path.join(path, f"{name}_history_main_mse.pdf")
PlotModelHistory(filename=filename, history=history, plot_metric="mse", main_branch=multiple_branches_used)
# plot learning rate
if lr_sc:
PlotModelLearningRate(filename=os.path.join(path, f"{name}_history_learning_rate.pdf"), lr_sc=lr_sc)