[Question] can `Scann` be used inside the model during training?

I have the following model using sequences to predict the next item:

class Model(tfrs.models.Model):
    def __init__(self):
        super().__init__()

        self.query_model = tf.keras.Sequential([
            QueryModel(),
            tf.keras.layers.Dense(64),
            L2NormalizationLayer(axis=1)

        ])
        
        self.candidate_model = tf.keras.Sequential([
            CandidateModel(),
            tf.keras.layers.Dense(64),
            L2NormalizationLayer(axis=1)
        ])
        
        scann = tfrs.layers.factorized_top_k.ScaNN(num_reordering_candidates=100)
        scann.index_from_dataset(
            candidates_ds.map(
                lambda x: (x['id'], self.candidate_model({ 'url': x['url'] }))
            )
        )
        
        self.task = tfrs.tasks.Retrieval(
            # Normal approach commented out, using the candidate model to map over a dataset of unique candidates.
             metrics=tfrs.metrics.FactorizedTopK(
                  # candidates=candidates_ds.map( lambda x: (x['id'], self.candidate_model({ 'url': x['url'] })))
                  # Instead we use the SCANN layer
                  candidates=scann
            )
            remove_accidental_hits=True
        )
        
    def call(self, features):
        candidate_embeddings = self.candidate_model({
            'url': features['url'],
        })

        query_embeddings = self.query_model({
            'advertiser_name': features['advertiser_name']        
        })
                
        return (
            query_embeddings,
            candidate_embeddings,
        )
    
    def compute_loss(self, features, training=False):
        query_embeddings, candidate_embeddings = self(features)

        return self.task(
            query_embeddings, 
            candidate_embeddings,
            candidate_ids=features['id'],
            compute_metrics=not training
        )

This runs fine and much quicker! The evaluation step sees 100x speed ups!

However this model does not improve on the metric, my first thought it that the model is not updating the embeddings each epoch as they are run only once. However, in the normal approach we also pass a dataset of already mapped candidate embeddings…

At which point in the training does the model update the embeddings it is learning to use for new evaluation runs?

[maciejkula]

@ydennisy to add to Patrick’s answer, Keras caches compiled TensorFlow functions. Remember to call compile before every evaluation as per https://github.com/tensorflow/recommenders/issues/388#issuecomment-941254103.

[patrickorlando]

This concept is discussed in https://github.com/tensorflow/recommenders/issues/388#issuecomment-941254103 and the comments following.

To make this work you must re-construct the index before each call to model.evaluate() to update the candidate embedddings.