Pipeline to obtain item embeddings of Amazon Review Data (2018).
Model is based on sequential recommender SASRec. [paper][github]
pip install -r src/requirements.txt
Our pipeline take .gz file from Amazon Review Data (2018) as input.
python data/data_process.py --file_path [gz file path] --output_path [output directory]
data_process.py will generate train/val/test json files with leave-one-out strategy as SASRec paper mentioned. "umap.json" and "smap.json" contain the dictionary that convert user id and item id in raw data into numbers used in model.
Please output all processed files under path data/dataset_name/
.
Name | Description | Example |
---|---|---|
train.json | Training set of item sequence, using the n-2 interactions | {0: [0, 1, 2, 3, 4]} |
val.json | Validation set of item sequence, using the (n-1)th interaction | {0: [5]} |
test.json | Testing set of item sequence, using the (n)th interaction | {0: [6]} |
umap.json | Mapping real user/sess id to consecutive integer id | {"A2HOI48JK8838M": 0} |
smap.json | Mapping real item id to consecutive integer id | {"B00004U9V2": 1} |
Model training with our default hyper-parameter configure.
- If GPU Mem is large enough to train the entire model
bash script/train.sh ${gpu_id} ${data_name} main
- If Multiple GPUs Mem are large enough to train the entire model
where we set
bash script/train_on_multi_devices.sh ${gpu_id} ${data_name}
--emb_device_idx "{'cuda:0':(0,20), 'cuda:1':(20,50)}"
to split embedding lookup table by dimension to feed into different devices. Note: we can also put partial embedding table (e.g. 20%) on cpu mem in this way, such as--emb_device_idx "{'cpu':(0,20), 'cuda:0':(20,50)}"
, but training will be super slow (~5x training time). - If GPU Mem is too small to train the entire model, and only one GPU, we use
small_gpu.py
to tackle the problem.Technical details: 1) We maintain a large table (e.g. 15M items) on CPU mem; 2) sample a subset (e.g. 5M items) as to GPU at one time for training, 3) then update the large table after training, 4) repeat the step 2) & 3) until convergence.bash script/train.sh ${gpu_id} ${data_name} small_gpu
Evaluation methods: NDCG@K, Recall@K, MRR, AUC
#Users | #Items | #Interactions | Sparsity |
---|---|---|---|
~44 Million | ~15 Million | ~0.2 Billion | 99.99997% |
For evaluation, we uniformly sample 100,000 users and for each user we uniformly sample 1000 items as negative candidates, which will be ranked with the single positive item (form a ranking list containing 1001 items).
Training time is about 25 hours on Large GPU
Mode and Multi GPU
Mode (with NVLink), and about 42 hours on Small GPU
Mode.
Training | NDCG@10 | Recall@10 | MRR | AUC |
---|---|---|---|---|
Large GPU Mode |
0.33591 | 0.50647 | 0.29566 | 0.89410 |
Small GPU Mode |
0.33855 | 0.499912 | 0.29800 | 0.80470 |
We also provide a notebook demo.ipynb
to demostrate:
- the whole training process with a toy Amazon dataset.
- how to export our model to
onnx
model and run withonnxruntime
ortensorrt
. - how to use our model to generate embeddings for approximated knn searching with
faiss
.