Training the retrieval model
With the towers built, the rest of training is short. Compile with AdamW, fit on the train_ds we set up earlier, and log the run to MLflow so we have an artifact we can serve later.
store = PostgresStore()
transactions = store.read_table('transactions', columns=['customer_id','article_id','t_dat','month_sin','month_cos'])
customers = store.read_table('customers', columns=['customer_id','age'])
articles = store.read_table('articles', columns=['article_id','garment_group_name','index_group_name'])
joined = (transactions.join(customers, on='customer_id').join(articles, on='article_id')).to_pandas()
dataset = TwoTowerDataset(joined, batch_size=settings.TWO_TOWER_MODEL_BATCH_SIZE)
train_ds, val_ds = dataset.get_train_val_split()
query_model = QueryTowerFactory(dataset).build()
item_model = ItemTowerFactory(dataset).build()
two_tower = TwoTowerFactory(dataset).build(query_model, item_model)
history = TwoTowerTrainer(dataset, two_tower).train(train_ds, val_ds)
registry.log_two_tower(query_model, item_model, metrics={'loss': history.history['loss'][-1]}, params={...})The training entrypoint reads features from Postgres, builds the dataset, trains both towers, and registers the artifacts.
On the sampled slice the loss curve is already plateauing by epoch four or five. Training longer mostly memorises noise. On the full Kaggle dataset you would push further, but you would also start checking val_loss every epoch and stopping when it diverges from train_loss. Watching the gap is more important than the epoch count.
Quiz: Quiz
Loading practice…