2025年5月23日更新：感谢 Wentworth1028 和 Tiny-Snow，我们引入了 LayerNorm 优化，显著提升了 NDCG 和 HR 指标。相关文档请见：doc 👍。

---

本项目基于论文作者的 TensorFlow 实现（GitHub 地址）进行修改，为简化代码迁移至 PyTorch (v1.6)，并修复了位置嵌入使用等问题（这有助于降低过拟合风险；不过在推荐系统中，个性化有时就意味着过拟合）。

代码位于 python 文件夹中。

训练命令如下：

python main.py --dataset=ml-1m --train_dir=default --maxlen=200 --dropout_rate=0.2 --device=cuda

仅进行推理的命令如下：

python main.py --device=cuda --dataset=ml-1m --train_dir=default --state_dict_path=[YOUR_CKPT_PATH] --inference_only=true --maxlen=200

由于负样本是随机采样的，每次运行的结果会略有差异。以下是连续两次运行的输出结果：

第1次运行 - 测试集 (NDCG@10: 0.5897, HR@10: 0.8190)
第2次运行 - 测试集 (NDCG@10: 0.5918, HR@10: 0.8225)

有关更详细的介绍和完整的 README，请参阅论文作者的仓库：GitHub 地址。此外，附上论文引用信息供参考：

@inproceedings{kang2018self,
  title={Self-attentive sequential recommendation},
  author={Kang, Wang-Cheng and McAuley, Julian},
  booktitle={2018 IEEE International Conference on Data Mining (ICDM)},
  pages={197--206},
  year={2018},
  organization={IEEE}
}

我注意到该仓库已被引用十余次🫰。如需引用，可参考以下示例：

@online{huang2020sasrec_pytorch,
  author  = {Zan Huang},
  title   = {SASRec.pytorch},
  year    = {2020},
  url     = {https://github.com/pmixer/SASRec.pytorch}
}

SASRec.pytorch 快速上手指南

SASRec.pytorch 是基于论文《Self-attentive sequential recommendation》的 PyTorch 实现版本。相较于原始的 TensorFlow 实现，该版本简化了代码结构，修复了位置嵌入（positional embedding）的使用问题，并支持 LayerNorm 以提升 NDCG 和 HR 指标。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：建议 Python 3.6+
• 深度学习框架：PyTorch v1.6 或更高版本
• 硬件加速：推荐使用 NVIDIA GPU 并安装对应的 CUDA 驱动（命令中需指定 --device=cuda），若无 GPU 可改为 cpu。

依赖安装： 建议使用国内镜像源加速安装 PyTorch 及其他依赖。

# 示例：使用清华源安装 PyTorch (请根据实际 CUDA 版本调整)
pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装其他可能需要的通用依赖
pip install numpy pandas tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目无需复杂的安装过程，直接克隆仓库即可使用。

git clone https://github.com/pmixer/SASRec.pytorch.git
cd SASRec.pytorch

所有核心代码均位于 python 文件夹内。

基本使用

1. 训练模型

使用 MovieLens-1M (ml-1m) 数据集进行训练。以下命令启用了 GPU 加速，设置了最大序列长度为 200，Dropout 率为 0.2。

python main.py --dataset=ml-1m --train_dir=default --maxlen=200 --dropout_rate=0.2 --device=cuda

注意：由于负采样具有随机性，每次运行的输出结果会有轻微波动。

2. 仅进行推理（测试）

如果您已有训练好的检查点文件（.ckpt），可以使用以下命令直接进行评估。请将 [YOUR_CKPT_PATH] 替换为您的实际模型路径。

python main.py --device=cuda --dataset=ml-1m --train_dir=default --state_dict_path=[YOUR_CKPT_PATH] --inference_only=true --maxlen=200

预期输出示例

运行结束后，终端将输出测试集的评估指标（如 NDCG@10 和 HR@10）。由于随机采样，数值会在一定范围内浮动：

1st run - test (NDCG@10: 0.5897, HR@10: 0.8190)
2nd run - test (NDCG@10: 0.5918, HR@10: 0.8225)

---

参考文献 如需引用此项目，请使用以下 BibTeX：

@online{huang2020sasrec_pytorch,
  author  = {Zan Huang},
  title   = {SASRec.pytorch},
  year    = {2020},
  url     = {https://github.com/pmixer/SASRec.pytorch}
}