GPT2

免责声明：这不是官方的GPT2实现！我已尽最大努力尽可能贴近原始GPT2模型的规范，但请注意，我尚未能用这段代码复现原始模型的全部性能。至于原因，我也不清楚，目前未能找到可能导致此问题的任何 bug。

这是一个支持 GPU 和 TPU 的 GPT2 训练实现。数据集脚本略显粗糙，可能需要根据您的需求进行调整。

需求

对于 GPU：

pip3 install tensorflow-gpu regex

对于 TPU：

pip3 install tensorflow regex google-api-python-client oauth2client

下载模型所需：

pip3 install requests tqdm

生成数据集所需（除 TensorFlow 外）：

pip3 install ftfy tqdm newspaper3k

下载预训练模型

如果您想使用我的模型，目前我提供了“117M”、“PrettyBig”和“1.5B”三种版本。“117M”是在单个 v2 TPU 上训练了一周（可能少于 OpenAI 原始模型的训练时间）；“PrettyBig”比 345M 略大，在 v2-256 Pod 上训练了一周。我原本还计划发布我的 1.5B 版本，但后来决定放弃。您可以阅读我的理由 这里。 由于 OpenAI 已经发布了他们的模型，我现在也发布了我的（质量稍逊的）1.5B 模型，该模型是在 v3-512 Pod 上训练了一周。

python3 download_model.py PrettyBig

这将创建两个目录，一个以模型命名，另一个名为“encoder”。请修改与您的模型对应的 .json 文件中的 “model_dir” 和 “encoder_path” 参数，分别指向这两个路径。

如果您只需要编码器，请使用：

python3 download_model.py encoder

文本生成

进行预测时，您可以直接在命令行中输入提示词，也可以从文件中读取提示词。（这对于包含换行符的提示词非常有用）文本会输出到控制台以及 “predict_path” 参数指定的文件中。为此，您需要将模型检查点和 BPE 编码器副本放置在可访问的位置。（请修改 .json 文件中的 “model_dir” 和 “encoder_path” 参数）

通过命令行：

python3 main.py --model Your-Model.json [--top_k Top-K-Truncation] --predict_text "Hello there! My name is"

从文件：

python3 main.py --model Your-Model.json [--top_k Top-K-Truncation] --predict_file input.txt

可选的 top_k 参数会使模型在每一步只考虑概率最高的前 k 个标记。将其设置为 40 左右通常会产生更好的结果，但多样性会降低。

不支持在 TPU 上进行预测。

训练

要训练一个模型，首先在 .json 文件中定义其参数（参见示例），然后直接运行：

python3 main.py --model Your-Model.json [--tpu Your-TPU-Name]

使用 TPU 是可选的，无需修改即可在 GPU 上正常运行。（注意：评估功能在 TPU Pod 上无法使用，必须注释掉）

此过程假定您已在可访问位置存储了 openwebtext 数据集。如果没有，请参阅下方如何生成您自己的版本。

生成数据集

GPT2 是基于 webtext 数据集训练的，该数据集基本上是 Reddit 上所有链接且至少有 3 个 Karma 的网站内容。由于数据库庞大且包含大量受版权保护的内容，我无法在此提供下载。相反，我将介绍我是如何获取它的。请注意，为了下载并处理整个数据集，我花费了大约 500 欧元的云资源费用，但我并不声称自己是最高效的。

1. 使用 此处 的下载脚本下载归档文件（我使用了预先过滤的 URL 文件）。
2. 使用 datasets/openwebtext/run_newspaper_extract.py 提取文本。
3. 获取原始 .txt 文件后，使用 datasets/openwebtext/create_tfrecords.py 将其编码为 .tfrecords 文件（需要编码器副本，详见“下载预训练模型”部分）。
4. 将 .tfrecords 文件放入可访问的文件夹或 Google Storage 存储桶中（如果使用 TPU，则必须放入 Google Storage 存储桶）。
5. 修改您的 .json 文件中的 “data_path” 参数，使其指向 .tfrecords 文件所在的位置；如有必要，还需调整 inputs.py 中的函数，以便在文件名更改时能够正确打开它们。

使用您自己的数据

您也可以使用自己的文本文件作为训练数据，但需要手动修改一些代码。

1. 修改 datasets/openwebtext/create_tfrecords.py 中的参数：

base_dir = "/home/connor/my_text_dir" # 您的 .txt 文件所在路径
files_per = 175000 # 每个 tfrecord 放入多少个 txt 文件，影响不大
name = "my-custom-data" # 输出文件的名称格式为 name_i.tfrecords，其中 i 是文件编号
output_dir = "/home/connor/output" # 保存 .tfrecords 文件的路径
log_dir = "logs" # 日志将存放于此，以便在编码中断时恢复
files = glob.glob(os.path.join(base_dir, "**/*.txt")) # 需要生成所有 txt 文件的路径列表
processes = 64 # 并行编码进程数
encoder_path = "/home/connor/encoder" # BPE 编码器文件路径
minimum_size = 128 # 文件允许的最小长度（以 BPE 令牌计），否则会被丢弃。

2. 运行脚本。这将生成一系列 name_i.tfrecords 文件。请将这些文件放置在一个可访问的位置（如果您使用 TPU，则必须放在 Google Cloud Storage 存储桶中）。
3. 在 inputs.py 中创建一个新的输入函数。任何输入函数都应具有签名 function_name(params, eval=False)。stitch 参数控制连接的文本数量，以避免样本过小。其值应为：ceil((n_ctx+1) / minimum_size)。例如，如果我的 minimum_size 是 128，而 n_ctx 是 1024，则 stitch 应为 9。

def my_input(params, eval=False):
    if not eval:
        numbers = [0, 3, 4, 5, 6, 7, 8, 9] # 训练用随机文件子集
    else:
        numbers = [1, 2] # 评估用随机文件子集
    files = [os.path.join(params["data_path"], "my-custom-data_{}.tfrecords".format(str(i))) for i in numbers] # 生成文件列表

    return bpe_text(params["batch_size"], files, amount=params["n_ctx"], iterations=params["iterations"], stitch=9, batch=True)

4. 在 main.py 中注册您的新输入。

inputs = {
    "openwebtext": openwebtext, # 标准 OpenWebText 输入
    "openwebtext_longbiased": openwebtext_longbiased, # 倾向于展示更多长文本（>512 个令牌）的 OpenWebText
    "openwebtext_long": openwebtext_long, # 只展示长文本的 OpenWebText
    "my_input": my_input,
}

5. 将您的 .json 配置设置为使用新输入。

[...]
    "iterations": 500,
    "n_embd": 768,
    "input": "my_input",
    "model": "GPT2",
[...]

6. 完成！此处描述的输入应尽可能接近 GPT2，并且可在 TPU 上完美运行。

参数说明

由于通过命令行传递二十多个参数会非常繁琐，因此所有模型参数都通过 .json 文件传递。请注意，所有路径也支持 Google Cloud Storage 路径；如果您在 TPU 上运行，则路径必须是 gs:// 格式。

您可能需要更改的值：

• model_path: 用于保存和加载检查点的路径
• data_path: 您的 .tfrecords 文件所在路径
• encoder_path: BPE 编码器文件的路径。您可以使用 download_model.py 脚本来下载任意模型（或仅下载编码器），得到一个名为“encoder”的文件夹，这就是您需要指向的路径（仅在预测时需要）。

您可能希望调整的值：

• train_batch_size: 训练阶段的批量大小
• eval_batch_size: 评估阶段的批量大小
• predict_batch_size: 预测阶段的批量大小
• predict_path: 用于保存预测结果的路径（可指向一个文本文件以追加内容）

模型参数：

• model: 指定使用的模型名称。此处应始终为“GPT2”（未实现其他模型）
• n_ctx: 模型一次处理的令牌数（默认：1024）
• n_vocab: 词汇表大小（默认：50257）
• n_embd: 嵌入层的维度
• n_layer: 模型的层数
• n_head: 注意力头的数量（默认：n_embd / 64）
• scale_by_depth: 是否根据层数缩放初始化权重（默认：是）
• scale_by_in: 是否根据输入通道数缩放初始化权重（默认：是）

训练参数：

• precision: 使用 float32 还是 bfloat16 变量（训练超大规模模型时使用“bfloat16”）（可选，默认为 float32）
• input: 使用哪个输入函数（默认：“openwebtext”）
• lr: 学习率（默认：0.00025）
• warmup_steps: 预热步数。若设置此参数，则采用线性预热 + 余弦退火调度（默认：2000）（可选）
• opt_name: 优化器名称，目前有“adam”和“adafactor”两种（默认：“adam”）
• weight_decay: 权重衰减参数，若未指定则不使用权重衰减（将应用 Adam 的权重衰减修正）（默认：0.01）（可选）
• beta1: Adam/Adafactor 的 beta1 参数（Adam 默认：0.9，Adafactor 默认：0.0）
• beta2: Adam/Adafactor 的 beta2 参数（默认：0.98）（对于采用幂次衰减类型的 Adafactor 可选）
• epsilon: Adam 的 epsilon 参数（默认：1e-9）
• decay_type: Adafactor 的衰减类型，可为“pow”或“adam”（默认：“pow”）
• decay_exponent: Adafactor 幂次衰减的指数（默认：0.8）
• train_steps: 两次评估之间的训练步数
• eval_steps: 每次评估的步数
• max_steps: 最大训练步数（对学习率衰减很重要）
• iterations: 在 TPU 上执行的迭代次数（默认：100）（仅适用于 TPU）
• embed_dropout: 词嵌入层的 dropout 率，设为 0 可禁用（默认：0.1）
• attn_dropout: 注意力层的 dropout 率，设为 0 可禁用（默认：0.1）
• res_dropout: 残差连接的 dropout 率，设为 0 可禁用（默认：0.1）

GPT2 快速上手指南

注意：本仓库并非 OpenAI 官方实现的 GPT-2。作者尽力还原了原始模型规格，但性能可能无法完全复现原版效果。

环境准备

本项目基于 TensorFlow，支持 GPU 和 TPU 训练。请根据您的硬件环境安装相应依赖。

系统要求

• Python 3.x
• TensorFlow (支持 GPU 或 TPU)

前置依赖安装

GPU 用户：

pip3 install tensorflow-gpu regex

TPU 用户：

pip3 install tensorflow regex google-api-python-client oauth2client

通用工具（下载模型及生成数据集所需）：

pip3 install requests tqdm ftfy newspaper3k

💡 国内加速建议：如遇下载缓慢，可指定国内镜像源安装，例如： pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow-gpu regex

---

安装与模型下载

本工具不提供预训练模型文件托管，需通过脚本下载作者提供的模型（包含 "117M", "PrettyBig", "1.5B" 等版本）或仅下载编码器。

下载预训练模型

以下命令将下载 "PrettyBig" 模型，并自动生成模型目录和 encoder 目录：

python3 download_model.py PrettyBig

仅下载编码器

如果您只需 BPE 编码器：

python3 download_model.py encoder

配置文件调整

下载完成后，请找到对应模型的 .json 配置文件，修改以下参数以指向实际路径：

• model_dir: 指向模型文件夹路径
• encoder_path: 指向 encoder 文件夹路径

---

基本使用

文本生成 (预测)

您可以直接通过命令行输入提示词，或从文件读取提示词进行文本生成。生成的文本将输出到控制台及配置文件中指定的路径。

注意：TPU 不支持预测功能，请在 GPU 环境下运行。

方式一：命令行直接输入

python3 main.py --model Your-Model.json [--top_k 40] --predict_text "Hello there! My name is"

注：--top_k 为可选参数，设置为 40 左右通常能生成质量更高但多样性稍低的文本。

方式二：从文件读取提示词

适用于包含换行符的长提示词：

python3 main.py --model Your-Model.json [--top_k 40] --predict_file input.txt

模型训练

训练前需准备好 .tfrecords 格式的数据集（默认使用 OpenWebText 或自定义数据），并在 .json 配置文件中设置好 data_path。

启动训练

python3 main.py --model Your-Model.json

使用 TPU 训练（可选）

python3 main.py --model Your-Model.json --tpu Your-TPU-Name

注：在 TPU Pod 上运行时，需在代码中注释掉评估（Evaluation）部分，因其暂不支持。

关键参数说明

所有模型与训练参数均在 .json 文件中配置，主要参数包括：

• model_path: 模型检查点保存/加载路径
• data_path: .tfrecords 数据集路径（TPU 用户必须使用 gs:// Google Storage 路径）
• n_ctx: 上下文窗口大小（默认 1024）
• lr: 学习率（默认 0.00025）
• train_batch_size: 训练批次大小