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