GPT-2 PyTorch 实现

GitHub

语言模型是无监督的多任务学习器

简介

本项目是 OpenAI GPT-2 模型的 PyTorch 实现。它提供了模型训练、句子生成和指标可视化功能。我们致力于使代码既易于理解又经过优化。我们设计的代码易于理解，同时我们还使用一些技术来提高性能。

依赖项

regex
tqdm
torch
numpy
matplotlib

使用方法

如何训练？

在训练 GPT-2 模型之前，需要准备语料库数据集。我们建议使用 Expanda 构建您自己的语料库。另外，训练模块需要分词后的训练和评估数据集以及它们的词汇表文件。

准备好数据集后，您可以按如下方式训练 GPT-2：

$ python -m gpt2 train --train_corpus           build/corpus.train.txt \
                       --eval_corpus            build/corpus.test.txt \
                       --vocab_path             build/vocab.txt \
                       --save_checkpoint_path   ckpt-gpt2.pth \
                       --save_model_path        gpt2-pretrained.pth
                       --batch_train            128 \
                       --batch_eval             128 \
                       --seq_len                64 \
                       --total_steps            1000000 \
                       --eval_steps             500 \
                       --save_steps             5000

要从上一个检查点文件恢复训练，请使用 --from_checkpoint [上一个检查点文件] 选项。如果您想使用多个 GPU 训练 GPT-2，请使用 --gpus [GPU 数量] 选项。

命令行用法详情如下：

usage: gpt2 train [-h] --train_corpus TRAIN_CORPUS --eval_corpus EVAL_CORPUS
                  --vocab_path VOCAB_PATH [--seq_len SEQ_LEN]
                  [--layers LAYERS] [--heads HEADS] [--dims DIMS]
                  [--rate RATE] [--dropout DROPOUT]
                  [--batch_train BATCH_TRAIN] [--batch_eval BATCH_EVAL]
                  [--base_lr BASE_LR] [--wd_rate WD_RATE]
                  [--total_steps TOTAL_STEPS] [--eval_steps EVAL_STEPS]
                  [--save_steps SAVE_STEPS]
                  [--save_model_path SAVE_MODEL_PATH]
                  [--save_checkpoint_path SAVE_CHECKPOINT_PATH]
                  [--from_checkpoint FROM_CHECKPOINT]
                  [--from_pretrained FROM_PRETRAINED] [--use_amp]
                  [--use_grad_ckpt] [--gpus GPUS]

可选参数:
  -h, --help            显示此帮助信息并退出

语料库和词汇表:
  --train_corpus TRAIN_CORPUS
                        训练语料库文件路径
  --eval_corpus EVAL_CORPUS
                        评估语料库文件路径
  --vocab_path VOCAB_PATH
                        词汇表文件路径

模型配置:
  --seq_len SEQ_LEN     最大序列长度
  --layers LAYERS       transformer 层数
  --heads HEADS         注意力层中的多头数量
  --dims DIMS           每层表示的维度
  --rate RATE           瓶颈层维度增加率
  --dropout DROPOUT     每个元素被丢弃的概率

训练和评估:
  --batch_train BATCH_TRAIN
                        训练批次大小
  --batch_eval BATCH_EVAL
                        评估批次大小
  --base_lr BASE_LR     默认学习率
  --wd_rate WD_RATE     权重衰减率
  --total_steps TOTAL_STEPS
                        总训练步数
  --eval_steps EVAL_STEPS
                        评估模型和记录指标的周期
  --save_steps SAVE_STEPS
                        将训练状态保存到检查点的周期

保存和恢复:
  --save_model_path SAVE_MODEL_PATH
                        将训练好的模型权重保存到文件
  --save_checkpoint_path SAVE_CHECKPOINT_PATH
                        将训练状态保存到检查点文件
  --from_checkpoint FROM_CHECKPOINT
                        从检查点文件加载上一次训练状态
  --from_pretrained FROM_PRETRAINED
                        从预训练模型初始化参数

扩展:
  --use_amp             在训练中使用自动混合精度
  --use_grad_ckpt       在 transformer 层中使用梯度检查点
  --gpus GPUS           训练中使用的 GPU 设备数量

生成句子！

训练 GPT-2 后，您可以在交互模式下使用训练好的模型生成句子。

$ python -m gpt2 generate --vocab_path      build/vocab.txt \
                          --model_path      model.pth \
                          --seq_len         64 \
                          --nucleus_prob    0.8

命令行用法详情如下：用法: gpt2 生成 [-h] --词表路径词表路径 --模型模型 [--序列长度序列长度] [--层数层数] [--头数头数] [--维度维度] [--比率比率] [--顶部概率顶部概率] [--使用GPU]

可选参数: -h, --帮助显示此帮助信息并退出 --词表路径词表路径词表文件路径 --模型路径模型路径训练好的GPT-2模型文件路径

模型配置: --序列长度序列长度最大序列长度 --层数层数 transformer层数 --头数头数注意力层中的多头数量 --维度维度每层表示的维度 --比率比率瓶颈层维度增长率

生成选项: --核采样概率核采样概率核采样的概率阈值 --使用GPU 在推理时使用GPU设备

评估模型

评估训练模型性能的一种方法是使用未在训练阶段使用的评估数据集计算客观指标。

$ python -m gpt2 evaluate --模型路径 model.pth --评估语料 corpus.test.txt --词表路径 vocab.txt

可视化指标

此外，您还可以通过可视化记录的指标来分析训练损失图。

$ python -m gpt2 visualize --模型路径 model.pth --交互式

示例图如下:

在训练中使用apex

在训练时，您可以使用NVIDIA apex来使用融合CUDA层和混合精度优化。选项--use_amp启用训练中的自动混合精度。在使用这些性能提升之前，您应该按照仓库的说明安装NVIDIA apex库，或运行以下命令:

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

如果您无法安装该库或您的GPU设备不支持快速混合精度训练(准确地说，GPU应支持通过Tensor Cores进行混合精度加速)，您可以以单精度模式训练模型。混合精度训练是一个可选项。在这种情况下，您仍然可以在训练中使用融合CUDA层，如Adam优化器和层归一化。

在Google Colab中体验!

您可以在Google Colab中体验训练好的GPT2模型!上面的notebook包含文本生成和指标评估。您需要将训练好的模型、词表文件和评估数据集上传到Google Cloud Storage。

对于对韩语版GPT2感兴趣的人，我们重写了上述notebook，特别提供了gpt2-ko-302M模型的案例，该模型使用约5.04B个来自韩语文档的标记进行训练。您可以在此notebook中体验演示。

许可证

本项目采用Apache-2.0许可证。

GPT2