def __init__( self, embedding_layer: tf.keras.layers.Layer, bottom_stack: Optional[tf.keras.layers.Layer] = None, feature_interaction: Optional[tf.keras.layers.Layer] = None, top_stack: Optional[tf.keras.layers.Layer] = None, task: Optional[tasks.Task] = None) -> None: """Initializes the model. Args: embedding_layer: The embedding layer is applied to categorical features. It expects a string-to-tensor (or SparseTensor/RaggedTensor) dict as an input, and outputs a dictionary of string-to-tensor of feature_name, embedded_value pairs. {feature_name_i: tensor_i} -> {feature_name_i: emb(tensor_i)}. bottom_stack: The `bottom_stack` layer is applied to dense features before feature interaction. If None, an MLP with layer sizes [256, 64, 16] is used. For DLRM model, the output of bottom_stack should be of shape (batch_size, embedding dimension). feature_interaction: Feature interaction layer is applied to the `bottom_stack` output and sparse feature embeddings. If it is None, DotInteraction layer is used. top_stack: The `top_stack` layer is applied to the `feature_interaction` output. The output of top_stack should be in the range [0, 1]. If it is None, MLP with layer sizes [512, 256, 1] is used. task: The task which the model should optimize for. Defaults to a `tfrs.tasks.Ranking` task with a binary cross-entropy loss, suitable for tasks like click prediction. """ super().__init__() self._embedding_layer = embedding_layer self._bottom_stack = bottom_stack if bottom_stack else layers.blocks.MLP( units=[256, 64, 16], final_activation="relu") self._top_stack = top_stack if top_stack else layers.blocks.MLP( units=[512, 256, 1], final_activation="sigmoid") self._feature_interaction = (feature_interaction if feature_interaction else feature_interaction_lib.DotInteraction()) if task is not None: self._task = task else: self._task = tasks.Ranking( loss=tf.keras.losses.BinaryCrossentropy( reduction=tf.keras.losses.Reduction.NONE ), metrics=[ tf.keras.metrics.AUC(name="auc"), tf.keras.metrics.BinaryAccuracy(name="accuracy"), ], prediction_metrics=[ tf.keras.metrics.Mean("prediction_mean"), ], label_metrics=[ tf.keras.metrics.Mean("label_mean") ] )
def __init__(self, candidate_dataset): super().__init__() self.query_model = tf.keras.layers.Dense(16) self.candidate_model = tf.keras.layers.Dense(16) self.ctr_model = tf.keras.layers.Dense(1, activation="sigmoid") self.retrieval_task = tasks.Retrieval( metrics=metrics.FactorizedTopK( candidates=candidate_dataset.map(self.candidate_model), ks=[5])) self.ctr_task = tasks.Ranking( metrics=[tf.keras.metrics.AUC(name="ctr_auc")])
def __init__(self, vocab_sizes: List[int], embedding_dim: int = 16, emb_optimizer: Optional[tf.keras.optimizers.Optimizer] = None, bottom_stack: Optional[tf.keras.layers.Layer] = None, feature_interaction: Optional[tf.keras.layers.Layer] = None, top_stack: Optional[tf.keras.layers.Layer] = None, task: Optional[tasks.Task] = None) -> None: super().__init__() emb_feature_config = _get_tpu_embedding_feature_config( vocab_sizes=vocab_sizes, embedding_dim=embedding_dim) if not emb_optimizer: emb_optimizer = tf.keras.optimizers.Adam() self._tpu_embeddings = embedding.TPUEmbedding(emb_feature_config, emb_optimizer) self._bottom_stack = bottom_stack if bottom_stack else MlpBlock( units=[256, 64, embedding_dim], out_activation="relu") self._top_stack = top_stack if top_stack else MlpBlock( units=[512, 256, 1], out_activation="sigmoid") self._feature_interaction = (feature_interaction if feature_interaction else DotInteraction()) if task is not None: self._task = task else: self._task = tasks.Ranking( loss=tf.keras.losses.BinaryCrossentropy( reduction=tf.keras.losses.Reduction.NONE), metrics=[ tf.keras.metrics.AUC(name="auc"), tf.keras.metrics.BinaryAccuracy(name="accuracy"), ], prediction_metrics=[ tf.keras.metrics.Mean("prediction_mean"), ], label_metrics=[tf.keras.metrics.Mean("label_mean")])
def __init__( self, embedding_layer: tf.keras.layers.Layer, bottom_stack: Optional[tf.keras.layers.Layer] = None, feature_interaction: Optional[tf.keras.layers.Layer] = None, top_stack: Optional[tf.keras.layers.Layer] = None, task: Optional[tasks.Task] = None) -> None: super().__init__() self._embedding_layer = embedding_layer self._bottom_stack = bottom_stack if bottom_stack else MlpBlock( units=[256, 64, 16], out_activation="relu") self._top_stack = top_stack if top_stack else MlpBlock( units=[512, 256, 1], out_activation="sigmoid") self._feature_interaction = (feature_interaction if feature_interaction else feature_interaction_lib.DotInteraction()) if task is not None: self._task = task else: self._task = tasks.Ranking( loss=tf.keras.losses.BinaryCrossentropy( reduction=tf.keras.losses.Reduction.NONE ), metrics=[ tf.keras.metrics.AUC(name="auc"), tf.keras.metrics.BinaryAccuracy(name="accuracy"), ], prediction_metrics=[ tf.keras.metrics.Mean("prediction_mean"), ], label_metrics=[ tf.keras.metrics.Mean("label_mean") ] )
def __init__(self): super().__init__() self._dense = tf.keras.layers.Dense(1) self.task = tasks.Ranking( loss=tf.keras.losses.BinaryCrossentropy(), metrics=[tf.keras.metrics.BinaryAccuracy(name="accuracy")])