# /*--------------------------------------------------------------------------------------------- # * Copyright (c) 2022-2023 STMicroelectronics. # * All rights reserved. # * # * This software is licensed under terms that can be found in the LICENSE file in # * the root directory of this software component. # * If no LICENSE file comes with this software, it is provided AS-IS. # *--------------------------------------------------------------------------------------------*/ import os import sys from timeit import default_timer as timer from datetime import timedelta from typing import Tuple, List, Dict, Optional import mlflow from hydra.core.hydra_config import HydraConfig from munch import DefaultMunch from omegaconf import DictConfig import numpy as np import tensorflow as tf def set_all_layers_trainable_parameter(model: tf.keras.Model, trainable=True) -> None: """ This function sets all elements in model.layers to be trainable or not depending on boolean value. Arguments: model (tf.keras.Model): the model. trainable (boolean): if True makes all layer trainable, if False makes all layers not trainable Returns: None """ model.trainable = trainable for layer in model.layers: layer.trainable = trainable def set_frozen_layers(model: tf.keras.Model, frozen_layers: str = None) -> None: """ This function freezes (makes non-trainable) the layers that are specified by the `frozen_layers` attribute. If `frozen_layers` is `None`, then all layers in the model are made trainable. The indices are specified using the same syntax as in Python: 2 layer 2 (3rd layer from network input) (2) layer 2 (4:8) layers 4 to 7 (0:-1) layers 0 to before last (7:13, 16, 20:28, -1) layers 7 to 12, 16, 20 to 27, last There can be any number of index slices in the attribute. The attribute must be between parentheses. Using square brackets would make it an array and it could be changed by the YAML loader. not be left unchanged by the Arguments: model (tf.keras.Model): the model. frozen_layers (str): the indices of the layers to freeze. If `None`, all layers are made trainable. Returns: None """ if frozen_layers == 'None': # Train the whole model set_all_layers_trainable_parameter(model, trainable=True) return frozen_layers = str(frozen_layers) frozen_layers = frozen_layers.replace(" ", "") frozen_slices = frozen_layers[1:-1] if frozen_layers[0] == "(" else frozen_layers num_layers = len(model.layers) indices = np.arange(num_layers) frozen_indices = np.zeros(num_layers, dtype=bool) for slice_ in frozen_slices.split(','): try: i = eval("indices[" + str(slice_) + "]") except ValueError as val_err: raise ValueError("\nInvalid syntax for `frozen_layers` attribute\nLayer index slices " f"should follow the Python syntax. Received {frozen_layers}\nPlease check the " "'training' section of your configuration file.") from val_err frozen_indices[i] = True # Freeze layers set_all_layers_trainable_parameter(model, trainable=True) num_trainable = num_layers for i in range(num_layers): if frozen_indices[i]: num_trainable -= 1 model.layers[i].trainable = False def set_dropout_rate(model: tf.keras.Model, dropout_rate: float = None) -> None: """ This function sets the dropout rate of the dropout layer if the `dropout`attribute was used in the 'training' section of the config file. An error is thrown if: - The `dropout` attribute was used but the model does not include a dropout layer. - The model includes a dropout layer but the `dropout`attribute was not used to specify a rate. Arguments: model (tf.keras.Model): the model. dropout_rate (float): the dropout rate to set on the dropout layer. Returns: None """ found_layer = False for layer in model.layers: layer_type = layer.__class__.__name__ if layer_type == "Dropout": found_layer = True if found_layer: # The dropout rate must be specified. if not dropout_rate: raise ValueError("\nThe model includes a dropout layer.\nPlease use the 'training.dropout' " "attribute in your configuration file to specify a dropout rate") for layer in model.layers: layer_type = layer.__class__.__name__ if layer_type == "Dropout": layer.rate = dropout_rate print("[INFO] : Set dropout rate to", dropout_rate) else: # The dropout rate can't be applied. if dropout_rate: raise ValueError("\nUnable to set the dropout rate specified by the 'training.dropout' " "attribute in the configuration file\nThe model does not include " "a dropout layer.") def get_optimizer(cfg: DictConfig) -> tf.keras.optimizers: """ This function creates a Keras optimizer from the 'optimizer' section of the config file. The optimizer name, attributes and values used in the config file are used to create a string that is the call to the Keras optimizer with its arguments. Then, the string is evaluated. If the evaluation succeeds, the optimizer object is returned. If it fails, an error is thrown with a message that tells the user that the name and/or arguments of the optimizer are incorrect. Arguments: cfg (DictConfig): dictionary containing the 'optimizer' section of the configuration file. Returns: tf.keras.optimizers: the Keras optimizer object. """ message = "\nPlease check the 'training.optimizer' section of your configuration file." if type(cfg) != DefaultMunch: raise ValueError(f"\nInvalid syntax for optimizer{message}") optimizer_name = list(cfg.keys())[0] optimizer_args = cfg[optimizer_name] # Get the optimizer if not optimizer_args: # The optimizer has no arguments. optimizer_text = f"tf.keras.optimizers.{optimizer_name}()" else: if type(optimizer_args) != DefaultMunch: raise ValueError(f"\nInvalid syntax for `{optimizer_name}` optimizer arguments{message}") text = f"tf.keras.optimizers.{optimizer_name}(" # Collect the arguments for k, v in optimizer_args.items(): if type(v) == str: text += f'{k}=r"{v}", ' else: text += f'{k}={v}, ' optimizer_text = text[:-2] + ')' try: optimizer = eval(optimizer_text) except ValueError as val_err: raise ValueError(f"\nThe optimizer name `{optimizer_name}` is unknown or," f"the arguments are invalid, got:\n{optimizer_text}.{message}") from val_err return optimizer