Transformer-XL：超越固定长度上下文的注意力语言模型

本仓库包含我们论文的 PyTorch 和 TensorFlow 两种实现代码：

Transformer-XL：超越固定长度上下文的注意力语言模型

Zihang Dai*, Zhilin Yang*, Yiming Yang, Jaime Carbonell, Quoc V. Le, Ruslan Salakhutdinov (*: 共同第一作者)

预印本 2018 年

TensorFlow

• 源代码位于 tf/ 文件夹中，支持 (1) 单节点多 GPU 训练，以及 (2) 多主机 TPU 训练。
• 除了源代码外，我们还提供了在论文中报告了最先进（SoTA）性能的预训练 “TensorFlow” 模型。
• 详情请参阅 tf/README.md。

PyTorch

• 源代码位于 pytorch/ 文件夹中，通过 nn.DataParallel 模块支持单节点多 GPU 训练。
• 详情请参阅 pytorch/README.md。

结果

Transformer-XL 在多个语言建模基准测试中取得了新的最先进结果。Transformer-XL 也是首个在字符级语言建模任务上突破 1.0 阈值的模型。以下是总结：

方法	enwiki8	text8	One Billion Word	WT-103	PTB（未微调）
先前最佳	1.06	1.13	23.7	20.5	55.5
Transformer-XL	0.99	1.08	21.8	18.3	54.5

致谢

getdata.sh 脚本的大部分内容来自 awd-lstm 仓库。祝大家语言建模愉快 :)

Transformer-XL 快速上手指南

Transformer-XL 是一种突破固定长度上下文限制的注意力语言模型，支持 PyTorch 和 TensorFlow 两种框架。本指南将帮助你快速配置环境并运行基础示例。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS
• GPU: 支持 CUDA 的 NVIDIA 显卡（可选，用于加速训练）
• Python: 3.6 或更高版本

前置依赖

请确保已安装以下基础工具：

• Git
• CUDA Toolkit (如需 GPU 加速)
• cuDNN (如需 GPU 加速)

建议国内用户使用清华源或阿里源加速 Python 包安装：

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目未打包为 PyPI 库，需克隆源码并使用。

1. 克隆仓库

git clone https://github.com/kimiyoung/transformer-xl.git
cd transformer-xl

2. 安装 Python 依赖

根据你选择的框架进入对应目录安装依赖。

PyTorch 用户:

cd pytorch
pip install -r requirements.txt

注：若 requirements.txt 中未指定 torch 版本，请手动安装适配你 CUDA 版本的 PyTorch：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

TensorFlow 用户:

cd tf
pip install -r requirements.txt

注：请确保安装的 TensorFlow 版本与代码兼容（通常推荐 TF 1.x 或特定 2.x 版本，具体参考 tf/requirements.txt）。

基本使用

以下以最常用的 PyTorch 版本为例，演示如何下载数据并启动训练。TensorFlow 用户请参考 tf/README.md 中的类似流程。

1. 准备数据集

运行提供的脚本下载并预处理标准数据集（如 WikiText-103）：

cd pytorch
bash getdata.sh

该脚本会自动下载数据至 data/ 目录。

2. 单卡训练示例

使用默认参数在 WikiText-103 数据集上开始训练：

python train.py --dataset wt103 --batch_size 32 --n_layer 12 --d_model 512 --n_head 8 --d_head 64 --d_inner 2048 --dropout 0.1 --dropatt 0.0

3. 多卡训练示例

利用 nn.DataParallel 进行单机多卡训练（例如使用 4 张 GPU）：

python train.py --dataset wt103 --batch_size 32 --n_layer 12 --d_model 512 --n_head 8 --d_head 64 --d_inner 2048 --dropout 0.1 --dropatt 0.0 --multi_gpu --gpu_ids 0,1,2,3

4. 评估模型

训练完成后，使用预训练权重或检查点进行困惑度（Perplexity）评估：

python eval.py --dataset wt103 --split test --batch_size 10 --clamp_len 400 --model_path best_model.pt

提示: 更多高级用法（如 TPU 训练、超参数调整）请分别查阅 pytorch/README.md 和 tf/README.md。