English | 中文

TencentPretrain: 腾讯预训练框架

预训练已成为人工智能技术的重要组成部分。TencentPretrain 是一个用于对不同模态数据（例如文本和视觉）进行预训练和微调的工具包。TencentPretrain 的特点是模块化设计，既便于使用现有的预训练模型，又为用户提供了进一步扩展的接口。借助 TencentPretrain，我们构建了一个包含多种性质预训练模型的模型库。TencentPretrain 继承了开源工具包 UER (https://github.com/dbiir/UER-py/)，并将其扩展为一个多模态预训练框架。

完整文档：https://github.com/Tencent/TencentPretrain/wiki

目录

• 特性
• 要求
• 快速入门
• 预训练数据
• 下游数据集
• 模型库
• 使用说明
• 竞赛解决方案
• 引用
• 联系方式

特性

TencentPretrain 具有以下特点：

• 可复现性 TencentPretrain 已在多个数据集上进行了测试，其性能应与原始预训练模型实现（如 BERT、GPT-2、ELMo、T5、CLIP）一致。
• 模型模块化 TencentPretrain 分为嵌入层、编码器、目标嵌入层（可选）、解码器（可选）和目标等部分。每个部分都实现了丰富的模块，清晰且稳健的接口使用户能够自由组合这些模块，以尽可能少的限制构建预训练模型。
• 多模态 TencentPretrain 支持文本、视觉和音频等多种模态。
• 模型训练 TencentPretrain 支持 CPU 模式、单 GPU 模式、分布式训练模式以及使用 DeepSpeed 进行超大规模模型训练。
• 模型库 在 TencentPretrain 的帮助下，我们预训练并发布了多种性质的模型。合理选择预训练模型对于下游任务的性能至关重要。
• SOTA 结果 TencentPretrain 支持全面的下游任务（如分类和机器阅读理解），并提供了多项竞赛的优胜方案。
• 丰富功能 TencentPretrain 提供了与预训练相关的丰富功能，例如特征提取器和文本生成等。

需求

• Python >= 3.6
• torch >= 1.1
• six >= 1.12.0
• argparse
• packaging
• regex
• 对于预训练模型的转换（与 TensorFlow 相关），您需要 TensorFlow
• 对于使用 SentencePiece 模型进行分词，您需要 SentencePiece
• 对于开发堆叠模型，您需要 LightGBM 和 BayesianOptimization
• 对于使用全词掩码进行预训练，您需要像 jieba 这样的分词工具
• 对于在序列标注下游任务中使用 CRF，您需要 pytorch-crf
• 对于超大规模模型训练，您需要 DeepSpeed
• 对于视觉模型训练，您需要 torchvision
• 对于音频模型训练，您需要 torchaudio，并且在某些 specaugment 的特殊设置中需要 opencv-python，在微调 speech2text 模型时则需要 editdistance

快速入门

本节通过几个常用示例展示如何使用 TencentPretrain。更多详细信息请参阅“使用说明”部分。我们首先在书评情感分类数据集上使用 BERT（一种文本预训练模型）。我们先在书评语料库上对模型进行预训练，然后再在书评情感分类数据集上进行微调。输入文件共有三个：书评语料库、书评情感分类数据集和词汇表。所有文件均采用 UTF-8 编码，并包含在本项目中。

BERT 语料库的格式如下（每行一个句子，文档之间用空行分隔）：

doc1-sent1
doc1-sent2
doc1-sent3

doc2-sent1

doc3-sent1
doc3-sent2

书评语料库是从书评情感分类数据集中获取的。我们移除标签，并从中间将每篇书评分成两部分，以构建一个包含两个句子的文档（详见 corpora 文件夹中的 book_review_bert.txt）。

分类数据集的格式如下：

label    text_a
1        instance1
0        instance2
1        instance3

标签和实例之间用制表符 \t 分隔。第一行为列名列表。对于 n 类分类问题，标签 ID 应为 0 到 n-1 之间的整数（包括 0 和 n-1）。

我们使用 Google 的中文词汇表文件 models/google_zh_vocab.txt，其中包含 21128 个汉字。

我们首先对书评语料库进行预处理。在预处理阶段，需要将语料库处理成指定预训练模型所需的格式（--data_processor）：

python3 preprocess.py --corpus_path corpora/book_review_bert.txt --vocab_path models/google_zh_vocab.txt \
                      --dataset_path dataset.pt --processes_num 8 --data_processor bert

请注意，需要安装 six>=1.12.0。

预处理过程较为耗时。使用多进程可以显著加快预处理速度（--processes_num）。默认使用 BERT 分词器（--tokenizer bert）。预处理完成后，原始文本会被转换为 dataset.pt，这是 pretrain.py 的输入。然后我们下载 Google 预训练的中文 BERT 模型 google_zh_model.bin（TencentPretrain 格式，原始模型来自 这里），并将其放入 models 文件夹中。我们加载该预训练的中文 BERT 模型，并在其基础上继续在书评语料库上进行预训练。预训练模型通常由嵌入层、编码器和目标层组成。为了构建预训练模型，我们需要提供相关信息。配置文件（--config_path）指定了预训练模型使用的模块和超参数。更多细节请参阅 models/bert/base_config.json。假设我们有一台配备 8 块 GPU 的机器：

python3 pretrain.py --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                    --pretrained_model_path models/google_zh_model.bin \
                    --config_path models/bert/base_config.json \
                    --output_model_path models/book_review_model.bin \
                    --world_size 8 --gpu_ranks 0 1 2 3 4 5 6 7 \
                    --total_steps 5000 --save_checkpoint_steps 1000 --batch_size 32

mv models/book_review_model.bin-5000 models/book_review_model.bin

请注意，由 pretrain.py 训练的模型会附加一个记录训练步数的后缀（--total_steps）。我们可以移除该后缀以便于使用。

接下来，我们在下游分类数据集上对预训练模型进行微调。我们使用 pretrain.py 输出的 book_review_model.bin 中的嵌入层和编码器：

python3 finetune/run_classifier.py --pretrained_model_path models/book_review_model.bin \
                                   --vocab_path models/google_zh_vocab.txt \
                                   --config_path models/bert/base_config.json \
                                   --train_path datasets/book_review/train.tsv \
                                   --dev_path datasets/book_review/dev.tsv \
                                   --test_path datasets/book_review/test.tsv \
                                   --epochs_num 3 --batch_size 32

微调后的分类模型默认保存路径为 models/finetuned_model.bin。需要注意的是，预训练任务的实际批量大小为 --batch_size 乘以 --world_size；而下游任务（如分类）的实际批量大小则为 --batch_size。

随后，我们使用微调后的模型进行推理：

python3 inference/run_classifier_infer.py --load_model_path models/finetuned_model.bin \
                                          --vocab_path models/google_zh_vocab.txt \
                                          --config_path models/bert/base_config.json \
                                          --test_path datasets/book_review/test_nolabel.tsv \
                                          --prediction_path datasets/book_review/prediction.tsv \
                                          --labels_num 2

--test_path 指定待预测文件的路径，该文件应包含 text_a 列。--prediction_path 指定包含预测结果的文件路径。我们需要通过 --labels_num 显式指定标签数量。上述数据集为二分类数据集。

以上内容提供了使用 TencentPretrain 进行预处理、预训练、微调和推理的基本方法。更多用法案例请参阅完整的 :arrow_right: 快速入门 :arrow_left: 。完整版快速入门包含了丰富的用例，覆盖了大多数与预训练相关的应用场景。建议用户阅读完整版快速入门，以便合理使用该项目。

预训练数据

本节提供了一系列 :arrow_right: 预训练数据 :arrow_left: 的链接。TencentPretrain 可以直接加载这些预训练数据。

下游数据集

本节提供了一系列 :arrow_right: 下游数据集 :arrow_left: 的链接。TencentPretrain 可以直接加载这些数据集。

模型库

借助 TencentPretrain，我们预训练了多种不同性质的模型（例如基于不同模态、编码器和目标的模型）。预训练模型的详细介绍及其下载链接请参阅 :arrow_right: 模型库 :arrow_left: 。所有预训练模型均可由 TencentPretrain 直接加载。

使用说明

TencentPretrain 的目录结构如下：

TencentPretrain/
    |--tencentpretrain/
    |    |--embeddings/ # 包含嵌入组件的模块
    |    |--encoders/ # 包含编码器组件的模块，如 RNN、CNN、Transformer
    |    |--decoders/ # 包含解码器组件的模块
    |    |--targets/ # 包含目标组件的模块，如语言模型、掩码语言模型
    |    |--layers/ # 包含常用的神经网络层
    |    |--models/ # 包含 model.py，用于组合不同组件的模块
    |    |--utils/ # 包含常用工具函数
    |    |--model_builder.py
    |    |--model_loader.py
    |    |--model_saver.py
    |    |--opts.py
    |    |--trainer.py
    |
    |--corpora/ # 包含预训练数据
    |--datasets/ # 包含下游任务数据集
    |--models/ # 包含预训练模型、词汇表和配置文件
    |--scripts/ # 包含用于预训练模型的实用脚本
    |--finetune/ # 包含下游任务微调脚本
    |--inference/ # 包含下游任务推理脚本
    |
    |--preprocess.py
    |--pretrain.py
    |--README.md
    |--README_ZH.md
    |--requirements.txt
    |--LICENSE

代码按照组件划分（如嵌入、编码器等），用户可以轻松使用并在此基础上进行扩展。

关于 TencentPretrain 的全面使用示例，请参阅 :arrow_right: 使用说明 :arrow_left: ，这些示例可以帮助用户快速实现 BERT、GPT-2、ELMo、T5、CLIP 等预训练模型，并在多种下游任务上对预训练模型进行微调。

竞赛解决方案

TencentPretrain 已被用于多项竞赛的获奖方案中。在本节中，我们提供了一些使用 TencentPretrain 在 CLUE 等竞赛中取得 SOTA 成绩的示例。更多详细信息请参阅 :arrow_right: 竞赛解决方案 :arrow_left: 。

引用

如果您在学术研究中使用了 TencentPretrain 中的工作（例如预训练模型），请引用发表于 ACL 2023 的 系统论文：

@article{zhao2023tencentpretrain,
  title={TencentPretrain: A Scalable and Flexible Toolkit for Pre-training Models of Different Modalities},
  author={Zhao, Zhe and Li, Yudong and Hou, Cheng and Zhao, Jing and others},
  journal={ACL 2023},
  pages={217},
  year={2023}
}

TencentPretrain 快速上手指南

TencentPretrain 是腾讯开源的多模态预训练框架，支持文本、视觉、音频等多种模态。它采用模块化设计，便于用户复现经典模型（如 BERT、GPT-2、CLIP 等）或构建自定义预训练模型。

环境准备

系统要求

• Python: >= 3.6
• PyTorch: >= 1.1
• 操作系统: Linux / macOS / Windows (推荐 Linux)

前置依赖

根据使用场景，可能需要安装以下额外依赖：

• 基础依赖: six, argparse, packaging, regex
• 分词工具: jieba (中文全词掩码预训练), SentencePiece
• 深度学习加速: DeepSpeed (超大模型训练), torchvision (视觉), torchaudio (音频)
• 其他: TensorFlow (模型转换), pytorch-crf (序列标注), LightGBM (堆叠模型)

建议先安装基础依赖：

pip install six>=1.12.0 argparse packaging regex

安装步骤

1. 克隆项目代码

   git clone https://github.com/Tencent/TencentPretrain.git
   cd TencentPretrain
   

2. 安装 Python 依赖 推荐使用国内镜像源加速安装：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注：若需使用 DeepSpeed 或特定模态功能，请参照“环境准备”章节单独安装对应库。

基本使用

以下示例演示如何使用 BERT 模型进行书评情感分类任务，流程包含：数据预处理 -> 继续预训练 -> 微调 -> 推理。

1. 数据预处理

将原始文本语料转换为框架所需的 .pt 格式。假设已有语料文件 corpora/book_review_bert.txt 和词表 models/google_zh_vocab.txt：

python3 preprocess.py --corpus_path corpora/book_review_bert.txt --vocab_path models/google_zh_vocab.txt \
                      --dataset_path dataset.pt --processes_num 8 --data_processor bert

提示：--processes_num 可指定多进程加速处理。

2. 继续预训练 (Optional)

加载谷歌官方中文 BERT 权重，在书评语料上继续预训练。假设机器有 8 张 GPU：

python3 pretrain.py --dataset_path dataset.pt --vocab_path models/google_zh_vocab.txt \
                    --pretrained_model_path models/google_zh_model.bin \
                    --config_path models/bert/base_config.json \
                    --output_model_path models/book_review_model.bin \
                    --world_size 8 --gpu_ranks 0 1 2 3 4 5 6 7 \
                    --total_steps 5000 --save_checkpoint_steps 1000 --batch_size 32

# 重命名模型文件（去除步数后缀）
mv models/book_review_model.bin-5000 models/book_review_model.bin

注：若无继续预训练需求，可直接使用下载的 google_zh_model.bin 进行微调。

3. 下游任务微调

使用预训练模型在标注好的分类数据集上进行微调：

python3 finetune/run_classifier.py --pretrained_model_path models/book_review_model.bin \
                                   --vocab_path models/google_zh_vocab.txt \
                                   --config_path models/bert/base_config.json \
                                   --train_path datasets/book_review/train.tsv \
                                   --dev_path datasets/book_review/dev.tsv \
                                   --test_path datasets/book_review/test.tsv \
                                   --epochs_num 3 --batch_size 32

微调后的模型默认保存为 models/finetuned_model.bin。

4. 模型推理

使用微调后的模型对无标签数据进行预测：

python3 inference/run_classifier_infer.py --load_model_path models/finetuned_model.bin \
                                          --vocab_path models/google_zh_vocab.txt \
                                          --config_path models/bert/base_config.json \
                                          --test_path datasets/book_review/test_nolabel.tsv \
                                          --prediction_path datasets/book_review/prediction.tsv \
                                          --labels_num 2

说明：--labels_num 需明确指定分类类别数量（此处为二分类）。

---

更多高级用法（如多模态训练、大模型分布式训练等）请参考官方完整文档：TencentPretrain Wiki