"""Training utilities""" # OS & env import os import logging import datetime import time # DS, ML & DL import numpy as np from sklearn.metrics import confusion_matrix, classification_report from keras.utils import image_dataset_from_directory from keras.layers import RandomFlip, RandomRotation, RandomZoom from keras.layers import GaussianNoise, RandomContrast, RandomBrightness from tensorflow.keras.callbacks import Callback, TensorBoard from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint import tensorflow as tf # images & data viz import matplotlib.pyplot as plt import seaborn as sns class ConditionalAugmentation(tf.keras.layers.Layer): def __init__(self, rate=0.2, **kwargs): super(ConditionalAugmentation, self).__init__(**kwargs) self.rate = rate self.flip = RandomFlip("horizontal") self.rotation = RandomRotation(0.25) self.zoom = RandomZoom(0.1) self.noise = GaussianNoise(0.1) self.contrast = RandomContrast(0.1) self.brightness = RandomBrightness(0.1) def call(self, inputs, training=None): if training: x = inputs x = tf.cond( tf.random.uniform(()) < self.rate, lambda: self.flip(x), lambda: x ) x = tf.cond( tf.random.uniform(()) < self.rate, lambda: self.rotation(x), lambda: x ) x = tf.cond( tf.random.uniform(()) < self.rate, lambda: self.zoom(x), lambda: x ) x = tf.cond( tf.random.uniform(()) < self.rate, lambda: self.noise(x), lambda: x ) x = tf.cond( tf.random.uniform(()) < self.rate, lambda: self.contrast(x), lambda: x ) x = tf.cond( tf.random.uniform(()) < self.rate, lambda: self.brightness(x), lambda: x ) return x return inputs def evaluate_model( model, model_arch, train_ds, val_ds, test_ds, LOG_DIR, CHKPT_DIR, model_name="raw_model", input_size=(224, 224), batch_size=32, n_epochs=10, optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy", "categorical_accuracy"], ) -> tuple: """Train, evaluate and log model from architecture and configuration Return model, history and plot confusion matrix """ if not os.path.exists(CHKPT_DIR): os.makedirs(CHKPT_DIR) chkpt_name = model_name + ".weights.h5" chkpt_uri = os.path.join(CHKPT_DIR, chkpt_name) model_config = f""" | Config | Value | |:---:|:---:| | **model name** | {model_name} | | **input size** | {input_size} | | **batch size** | {batch_size} | | **n epochs** | {n_epochs} | | **optimizer** | {optimizer} | | **loss** | {loss} | | **metrics** | {metrics} | | **best weights URI** | {chkpt_uri} | """ # set log folder log_dir = os.path.join( LOG_DIR, model_name, datetime.datetime.now().strftime("%Y%m%d-%H%M%S") ) # COMPLIE logging.info("⚙️ compiling") model.compile( optimizer=optimizer, loss=loss, metrics=metrics, ) # CALLBACKS logging.info("🛎️ declaring callbacks") class TimingCallback(Callback): def __init__(self): self.logs = [] self.start_time = None def on_train_begin(self, logs={}): self.start_time = time.time() # log time by epoch def on_epoch_end(self, epoch, logs={}): self.logs.append(time.time() - self.start_time) # log total time def on_train_end(self, logs={}): self.tot_time_sec = time.time() - self.start_time self.total_time = f"Total train time: {self.tot_time_sec // 60 :.0f}'{self.tot_time_sec % 60 :.0f}s" timing_callback = TimingCallback() checkpoint = ModelCheckpoint( chkpt_uri, save_best_only=True, save_weights_only=True, ) early_stopping = EarlyStopping( monitor="val_loss", patience=6, restore_best_weights=True ) tensorboard_callback = TensorBoard( log_dir=log_dir, histogram_freq=0, # do not save weights & biases (too much memory) write_graph=True, write_images=True, update_freq="epoch", ) # FIT logging.info("💪 starting training") model_history = model.fit( train_ds, validation_data=val_ds, epochs=n_epochs, callbacks=[timing_callback, checkpoint, early_stopping, tensorboard_callback], ) # EVALUATE ON TEST DATASET logging.info("🧐 evaluating model") model.load_weights(chkpt_uri) test_loss, *test_metrics = model.evaluate(test_ds) predictions = model.predict(test_ds) # CONFUSION MATRIX logging.info("📈 plotting results") # get true labels from test dataset true_labels = np.concatenate([y for x, y in test_ds], axis=0) # convert predictions to classes predicted_classes = np.argmax(predictions, axis=1) # compute confusion matrix conf_matrix = confusion_matrix(true_labels, predicted_classes) # precision & F1 score report = classification_report( true_labels, predicted_classes, target_names=test_ds.class_names, ) report_dict = classification_report( true_labels, predicted_classes, target_names=test_ds.class_names, output_dict=True, ) print(report) # plot it conf_mtx_plot = plt.figure(figsize=(6, 4)) sns.heatmap( conf_matrix, annot=True, fmt="d", cmap="Blues", xticklabels=test_ds.class_names, yticklabels=test_ds.class_names, ) plt.suptitle(f"{model_name} model", color="blue", weight="bold") plt.title( f"acc. {report_dict['accuracy'] :.02f} - loss {test_loss :.02f} - {timing_callback.total_time}", fontsize=10, ) plt.xlabel("Predictions", color="red", weight="bold") plt.ylabel("True labels", color="green", weight="bold") plt.show() # convert image for Tensorboard conf_mtx_plot.canvas.draw() image_array = np.array(conf_mtx_plot.canvas.renderer.buffer_rgba()) conf_mtx_plot_tf = tf.convert_to_tensor(image_array) conf_mtx_plot_tf = tf.expand_dims(conf_mtx_plot_tf, 0) plt.close() # LOG IN TENSORBOARD logging.info("📓 logging results") file_writer = tf.summary.create_file_writer(log_dir + "/metrics") with file_writer.as_default(): tf.summary.text("configuration", model_config, step=0) tf.summary.text("architecture", model_arch, step=0) tf.summary.text("total_training_time", timing_callback.total_time, step=0) for i, time_per_epoch in enumerate(timing_callback.logs): tf.summary.scalar("time_per_epoch", time_per_epoch, step=i + 1) tf.summary.image("confusion_matrix", conf_mtx_plot_tf, step=0) return model, model_history def eval_pretrained_model( model, train_ds, val_ds, test_ds, LOG_DIR, CHKPT_DIR, model_name="raw_model", input_size=(224, 224), batch_size=32, n_epochs=10, optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"], ) -> tuple: """Train, evaluate and log pre-trained model from architecture and configuration Return model, history and plot confusion matrix """ if not os.path.exists(CHKPT_DIR): os.makedirs(CHKPT_DIR) chkpt_name = model_name + ".weights.h5" chkpt_uri = os.path.join(CHKPT_DIR, chkpt_name) model_config = f""" | Config | Value | |:---:|:---:| | **model name** | {model_name} | | **input size** | {input_size} | | **batch size** | {batch_size} | | **n epochs** | {n_epochs} | | **optimizer** | {optimizer} | | **loss** | {loss} | | **metrics** | {metrics} | | **best weights URI** | {chkpt_uri} | """ # set log folder log_dir = os.path.join( LOG_DIR, model_name, datetime.datetime.now().strftime("%Y%m%d-%H%M%S") ) # COMPLIE logging.info("⚙️ compiling") model.compile( optimizer=optimizer, loss=loss, metrics=metrics, ) # CALLBACKS logging.info("🛎️ declaring callbacks") class TimingCallback(Callback): def __init__(self): self.logs = [] self.start_time = None def on_train_begin(self, logs={}): self.start_time = time.time() # log time by epoch def on_epoch_end(self, epoch, logs={}): self.logs.append(time.time() - self.start_time) # log total time def on_train_end(self, logs={}): self.tot_time_sec = time.time() - self.start_time self.total_time = f"Total train time: {self.tot_time_sec // 60 :.0f}'{self.tot_time_sec % 60 :.0f}s" timing_callback = TimingCallback() checkpoint = ModelCheckpoint( chkpt_uri, save_best_only=True, save_weights_only=True, ) early_stopping = EarlyStopping( monitor="val_loss", patience=10, restore_best_weights=True ) tensorboard_callback = TensorBoard( log_dir=log_dir, histogram_freq=0, # do not save weights & biases (too much memory) write_graph=True, write_images=True, update_freq="epoch", ) # FIT logging.info("💪 starting training") model_history = model.fit( train_ds, validation_data=val_ds, epochs=n_epochs, callbacks=[timing_callback, checkpoint, early_stopping, tensorboard_callback], ) # EVALUATE ON TEST DATASET logging.info("🧐 evaluating model") model.load_weights(chkpt_uri) test_loss, *test_metrics = model.evaluate(test_ds) predictions = model.predict(test_ds) # CONFUSION MATRIX logging.info("📈 plotting results") # get true labels from test dataset true_labels = np.concatenate([y for x, y in test_ds], axis=0) # convert predictions to classes predicted_classes = np.argmax(predictions, axis=1) # compute confusion matrix conf_matrix = confusion_matrix(true_labels, predicted_classes) # precision & F1 score report = classification_report( true_labels, predicted_classes, target_names=test_ds.class_names, ) report_dict = classification_report( true_labels, predicted_classes, target_names=test_ds.class_names, output_dict=True, ) print(report) # plot it conf_mtx_plot = plt.figure(figsize=(6, 4)) sns.heatmap( conf_matrix, annot=True, fmt="d", cmap="Blues", xticklabels=test_ds.class_names, yticklabels=test_ds.class_names, ) plt.suptitle(f"{model_name} model", color="blue", weight="bold") plt.title( f"acc. {report_dict['accuracy'] :.02f} - loss {test_loss :.02f} - {timing_callback.total_time}", fontsize=10, ) plt.xlabel("Predictions", color="red", weight="bold") plt.ylabel("True labels", color="green", weight="bold") plt.show() # convert image for Tensorboard conf_mtx_plot.canvas.draw() image_array = np.array(conf_mtx_plot.canvas.renderer.buffer_rgba()) conf_mtx_plot_tf = tf.convert_to_tensor(image_array) conf_mtx_plot_tf = tf.expand_dims(conf_mtx_plot_tf, 0) plt.close() # LOG IN TENSORBOARD logging.info("📓 logging results") file_writer = tf.summary.create_file_writer(log_dir + "/metrics") with file_writer.as_default(): tf.summary.text("configuration", model_config, step=0) tf.summary.text("total_training_time", timing_callback.total_time, step=0) for i, time_per_epoch in enumerate(timing_callback.logs): tf.summary.scalar("time_per_epoch", time_per_epoch, step=i + 1) tf.summary.image("confusion_matrix", conf_mtx_plot_tf, step=0) return model, model_history if __name__ == "__main__": help()