Added DLRM notebook and generated python script #2131
Conversation
There was a problem hiding this comment.
Thanks for the PR! Left some comments, let me know what you think :).
Here are some notes for posterity (we don't need to implement these now in this PR):
- We should, in the future, also think about the specific case of how how DLRM implements distributed training. The embeddings are sharded across devices, whereas the MLP layers are DDP'd.
- The paper mentions that they use
torch.nn.EmbeddingBaginstead oftorch.nn.Embedding. Apparently,torch.nn.EmbeddingBagis more efficient thantorch.nn.Embeddingfollowed by some aggregation op. We should think about this at some later point (maybe,keras_rs.embeddings.DistributedEmbeddingsolves this). Anyway, we can worry about this later!
Updated docstrings to match the context according to code
|
@kharshith-k - let me know when this is ready for another round of review. Thanks! |
|
Delete duplicate dlrm.ipynb
|
Hi @abheesht17, The changes are ready for review. I've used raw_user_age and timestamp for continuous features. Please let me know if any more changes would be required. Thanks! |
| normalization_layers = {} | ||
| for feature_name in MOVIELENS_CONFIG["continuous_features"]: | ||
| normalization_layers[feature_name] = keras.layers.Normalization(axis=-1) |
There was a problem hiding this comment.
Sorry to be a pain here, but could you extract the hour of day, hour of week from the timestamp, and map it to a unit circle using sin, cos? That's a better way of encoding timestamp
There was a problem hiding this comment.
Sure not a problem! I will do the changes accordingly
examples/keras_rs/dlrm.py
Outdated
| # Layers for continuous features. | ||
| self.continuous_dense_layers = {} | ||
| for feature_name in MOVIELENS_CONFIG["continuous_features"]: | ||
| # Use a small MLP to process continuous features | ||
| self.continuous_dense_layers[feature_name] = keras.Sequential( | ||
| [ | ||
| keras.layers.Dense(mlp_dim, activation="relu"), | ||
| keras.layers.Dense(mlp_dim), | ||
| ] | ||
| ) |
There was a problem hiding this comment.
Generally, we concatenate all dense features together, and use just one MLP. Let's do that, instead of creating separate MLPs from every dense feature.
|
Hello @abheesht17 I've modified the DLRM code as per your suggestions. I've encoded the timestamp feature by splitting it into hour_of_day_sin, hour_of_day_cos, hour_of_week_sin and hour_of_week_cos using sin and cos functions in the preprocessing. Please let me know if its correct or needs any corrections. Thanks! |
Added DLRM.ipynb and generated dlrm.py