Skip to content
Source code and datasets for the paper "Session-based Recommendation with Graph Neural Networks" (AAAI-19)
Branch: master
Clone or download
Latest commit aa51065 May 17, 2019
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
datasets Update preprocess.py Jan 22, 2019
pytorch_code SR-GNN Nov 12, 2018
tensorflow_code SR-GNN Nov 12, 2018
README.md Update README.md May 17, 2019

README.md

SR-GNN

Paper data and code

This is the code for the AAAI 2019 Paper: Session-based Recommendation with Graph Neural Networks. We have implemented our methods in both Tensorflow and Pytorch.

Here are two datasets we used in our paper. After downloaded the datasets, you can put them in the folder datasets/:

There is a small dataset sample included in the folder datasets/, which can be used to test the correctness of the code.

We have also written a blog explaining the paper.

Usage

You need to run the file datasets/preprocess.py first to preprocess the data.

For example: cd datasets; python preprocess.py --dataset=sample

usage: preprocess.py [-h] [--dataset DATASET]

optional arguments:
  -h, --help         show this help message and exit
  --dataset DATASET  dataset name: diginetica/yoochoose/sample

Then you can run the file pytorch_code/main.py or tensorflow_code/main.py to train the model.

For example: cd pytorch_code; python main.py --dataset=sample

You can add the suffix --nonhybrid to use the global preference of a session graph to recommend instead of the hybrid preference.

You can also change other parameters according to the usage:

usage: main.py [-h] [--dataset DATASET] [--batchSize BATCHSIZE]
               [--hiddenSize HIDDENSIZE] [--epoch EPOCH] [--lr LR]
               [--lr_dc LR_DC] [--lr_dc_step LR_DC_STEP] [--l2 L2]
               [--step STEP] [--patience PATIENCE] [--nonhybrid]
               [--validation] [--valid_portion VALID_PORTION]

optional arguments:
  -h, --help            show this help message and exit
  --dataset DATASET     dataset name:
                        diginetica/yoochoose1_4/yoochoose1_64/sample
  --batchSize BATCHSIZE
                        input batch size
  --hiddenSize HIDDENSIZE
                        hidden state size
  --epoch EPOCH         the number of epochs to train for
  --lr LR               learning rate
  --lr_dc LR_DC         learning rate decay rate
  --lr_dc_step LR_DC_STEP
                        the number of epochs after which the learning rate
                        decay
  --l2 L2               l2 penalty
  --step STEP           gnn propogation steps
  --patience PATIENCE   the number of epoch to wait before early stop
  --nonhybrid           only use the global preference to predict
  --validation          validation
  --valid_portion VALID_PORTION
                        split the portion of training set as validation set

Requirements

  • Python 3
  • PyTorch 0.4.0 or Tensorflow 1.9.0

Other Implementation for Reference

There are other implementation available for reference:

  • Implementation based on PaddlePaddle by Baidu [Link]
  • Implementation based on PyTorch Geometric [Link]
  • Another implementation based on Tensorflow [Link]
  • Yet another implementation based on Tensorflow [Link]

Citation

Please cite our paper if you use the code:

@inproceedings{Wu:2019vb,
    author = {Wu, Shu and Tang, Yuyuan and Zhu, Yanqiao and Wang, Liang and Xie, Xing and Tan, Tieniu},
	title = {Session-based Recommendation with Graph Neural Networks},
	booktitle = {Proceedings of The Twenty-Third AAAI Conference on Artificial Intelligence},
	series = {AAAI '19},
	year = {2019},
	url = {http://arxiv.org/abs/1811.00855}
}
You can’t perform that action at this time.