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

[ydennisy]

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?

[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.

[AndrewJGroves]

Sorry, just to confirm. Does this mean you can't use Scann when you are fitting your model for the first time? (Trying to speed up my implemention)

[patrickorlando]

ScaNN is only used for efficient retrieval. It can only help with speeding up evaluation and has no impact on training step speed.

[AndrewJGroves]

Thanks for the reply, is there any suggestions on how to speed up the retrieval task (bar using GPUs)

[patrickorlando]

There's not much specific advice I can give if you are running on CPU only, other than the best practices around using the tf.data.Dataset API with parallelism. Moving your lookups into the data pipeline and checking out the Tensorboard Profiler to identify bottlenecks.

[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.