Skip to content

oceanlvr/ProtoAU

Repository files navigation

ProtoAU

Pytorch implementation of ProtoAU for recomandation. We present the Prototypical contrastive learning through Alignment and Uniformity for recommendation, which is called ProtoAU. A contrastive learning method for recommendation that excels in capturing intricate relationships between user and item interactions, which enhance the basic GNN-based recommendation model's generalization ability and robustness.

Thanks for following our work! :)

Prepare

There are two environment you can choose: nvidia-docker environment or normal environment.

  • For nvidia-docker users, you need to install nvidia-docker2 and restart docker service.
# docker env
docker build -t protoau .
docker run -itd --gpus all --name protoau
docker exec -it protoau /bin/bash # enter the container
  • For normal users, you need to install pytorch and other packages. here we use follow environment:
    • Python 3.6
    • Pytorch 1.9 (GPU version)
    • CUDA 11.1
    • cudnn 8

then run follow command to install other packages:

pip install -r requirements.txt

Quickstart

  • Arguments:

    • Config the model arguments in conf/ProtoAU.yaml
  • Train:

# train
nohup python index.py --gpu_id=0 --model=ProtoAU --run_name=ProtoAU --dataset=yelp2018 > ./0.log 2>&1 &

# Parallel train(optional)
wandb sweep --project sweep_parallel ./sweep/ProtoAU.yaml # step 1
wandb agent --count 5 oceanlvr/sweep_parallel/[xxx] # replace the [xxx] with your sweep id (step 1 generated)
  1. For all metric results, you could see the output in the ./0.log file or the wandb dashboard.
  2. For visualizing the results, run python3 visualize/feature.py.

Datasets

   
DataSet
Users ItemsRatings Density
Douban 2,848 39,586 894,887 0.794%
LastFM 1,892 17,632 92,834 0.27%
Yelp 19,539 21,266 450,884 0.11%
Amazon-Book 52,463 91,599 2,984,108 0.11%

Reference

Cite

Please cite our paper (arXiv:2402.02079) if you use this code.

Releases

No releases published

Packages

No packages published