def __init__( self, feature_config: FeatureConfig, tfrecord_type: str, file_io: FileIO, scorer: Optional[ScorerBase] = None, metrics: List[Union[Type[kmetrics.Metric], str]] = [], optimizer: Optional[Optimizer] = None, model_file: Optional[str] = None, initialize_layers_dict: dict = {}, freeze_layers_list: list = [], compile_keras_model: bool = False, output_name: str = "score", logger=None, ): """ Constructor to instantiate a RelevanceModel that can be used for training and evaluating the search ML task Parameters ---------- feature_config : `FeatureConfig` object FeatureConfig object that defines the features to be loaded in the dataset and the preprocessing functions to be applied to each of them tfrecord_type : {"example", "sequence_example"} Type of the TFRecord protobuf message used for TFRecordDataset file_io : `FileIO` object file I/O handler objects for reading and writing data scorer : `ScorerBase` object Scorer object that wraps an InteractionModel and converts input features into scores metrics : list List of keras Metric classes that will be used for evaluating the trained model optimizer : `Optimizer` Tensorflow keras optimizer to be used for training the model model_file : str, optional Path to pretrained model file to be loaded for evaluation or retraining initialize_layers_dict : dict, optional Dictionary of tensorflow layer names mapped to the path of pretrained weights Use this for transfer learning with pretrained weights freeze_layers_list : list, optional List of model layer names to be frozen Use this for freezing pretrained weights from other ml4ir models compile_keras_model : bool, optional Whether the keras model loaded from disk should be compiled with loss, metrics and an optimizer output_name : str, optional Name of the output tensorflow node that captures the score logger : `Logger`, optional logging handler for status messages """ self.feature_config: FeatureConfig = feature_config self.logger: Logger = logger self.output_name = output_name self.scorer = scorer self.tfrecord_type = tfrecord_type self.file_io = file_io if scorer: self.max_sequence_size = scorer.interaction_model.max_sequence_size else: self.max_sequence_size = 0 # Load/Build Model if model_file and not compile_keras_model: """ If a model file is specified, load it without compiling into a keras model NOTE: This will allow the model to be only used for inference and cannot be used for retraining. """ self.model: Model = self.load(model_file) self.is_compiled = False else: """ Specify inputs to the model Individual input nodes are defined for each feature Each data point represents features for all records in a single query """ inputs: Dict[str, Input] = feature_config.define_inputs() scores, train_features, metadata_features = scorer(inputs) # Create model with functional Keras API self.model = Model(inputs=inputs, outputs={self.output_name: scores}) self.model.output_names = [self.output_name] # Get loss fn loss_fn = scorer.loss.get_loss_fn(**metadata_features) # Get metric objects metrics_impl: List[Union[str, kmetrics.Metric]] = get_metrics_impl( metrics=metrics, feature_config=feature_config, metadata_features=metadata_features) # Compile model """ NOTE: Related Github issue: https://github.com/tensorflow/probability/issues/519 """ self.model.compile( optimizer=optimizer, loss=loss_fn, metrics=metrics_impl, experimental_run_tf_function=False, ) # Write model summary to logs model_summary = list() self.model.summary(print_fn=lambda x: model_summary.append(x)) if self.logger: self.logger.info("\n".join(model_summary)) if model_file: """ If model file is specified, load the weights from the SavedModel NOTE: The architecture, loss and metrics of self.model need to be the same as the loaded SavedModel """ self.load_weights(model_file) # Initialize layer weights for layer_name, layer_file in initialize_layers_dict.items(): layer = self.model.get_layer(layer_name) layer.set_weights( self.file_io.load_numpy_array(layer_file, unzip=True)) self.logger.info("Setting {} weights from {}".format( layer_name, layer_file)) # Freeze layer weights for layer_name in freeze_layers_list: layer = self.model.get_layer(layer_name) layer.trainable = False self.logger.info("Freezing {} layer".format(layer_name)) self.is_compiled = True
def __init__( self, feature_config: FeatureConfig, tfrecord_type: str, file_io: FileIO, scorer: Optional[ScorerBase] = None, metrics: List[Union[Type[kmetrics.Metric], str]] = [], optimizer: Optional[Optimizer] = None, model_file: Optional[str] = None, compile_keras_model: bool = False, output_name: str = "score", logger=None, ): """Use this constructor to define a custom scorer""" self.feature_config: FeatureConfig = feature_config self.logger: Logger = logger self.output_name = output_name self.scorer = scorer self.tfrecord_type = tfrecord_type self.file_io = file_io if scorer: self.max_sequence_size = scorer.interaction_model.max_sequence_size else: self.max_sequence_size = 0 # Load/Build Model if model_file and not compile_keras_model: """ If a model file is specified, load it without compiling into a keras model NOTE: This will allow the model to be only used for inference and cannot be used for retraining. """ self.model: Model = self.load(model_file) self.is_compiled = False else: """ Specify inputs to the model Individual input nodes are defined for each feature Each data point represents features for all records in a single query """ inputs: Dict[str, Input] = feature_config.define_inputs() scores, train_features, metadata_features = scorer(inputs) # Create model with functional Keras API self.model = Model(inputs=inputs, outputs={self.output_name: scores}) # Get loss fn loss_fn = scorer.loss.get_loss_fn(**metadata_features) # Get metric objects metrics_impl: List[Union[str, kmetrics.Metric]] = get_metrics_impl( metrics=metrics, feature_config=feature_config, metadata_features=metadata_features ) # Compile model """ NOTE: Related Github issue: https://github.com/tensorflow/probability/issues/519 """ self.model.compile( optimizer=optimizer, loss=loss_fn, metrics=metrics_impl, experimental_run_tf_function=False, ) # Write model summary to logs model_summary = list() self.model.summary(print_fn=lambda x: model_summary.append(x)) self.logger.info("\n".join(model_summary)) if model_file: """ If model file is specified, load the weights from the SavedModel NOTE: The architecture, loss and metrics of self.model need to be the same as the loaded SavedModel """ self.load_weights(model_file) self.is_compiled = True