如何在Python中使用Transformers从头开始​​训练BERT？

2021年11月11日02:08:50 发表评论 2,190 次浏览

了解如何使用 Python 中的 Huggingface Transformers 库在自定义数据集上的掩码语言建模 (MLM) 任务中预训练 BERT 和其他转换器。

Python使用Transformers训练BERT？预训练模型是先前在大型数据集上训练并保存以供直接使用或微调的模型。在本教程中，你将学习如何借助 Python 中的 Huggingface 转换器库在自定义文本数据集上预训练 BERT（或任何其他转换器模型），包括Transformers训练BERT示例。

Transformers 的预训练可以通过自监督任务完成，以下是一些在 BERT 上完成的流行任务：

  • 掩码语言建模（MLM）：该任务包括对句子中一定比例的标记进行掩码，并训练模型来预测这些掩码单词。我们将在本教程中使用这个。
  • Next Sentence Prediction (NSP)：模型接收成对的句子作为输入，并学习预测成对中的第二个句子是否是原始文档中的后续句子。

首先，我们需要安装 3 个库：

$ pip install datasets transformers==4.11.2 sentencepiece

如果你想继续，请打开一个新的笔记本或 Python 文件并导入必要的库：

from datasets import *
from transformers import *
from tokenizers import *
import os
import json

选择数据集

Transformers如何训练BERT？如果你愿意对转换器进行预训练，那么你很可能拥有自定义数据集。但出于本教程中的演示目的，我们将使用cc_news数据集，为此我们将使用Huggingface 数据集库。因此，请确保按照此链接将你的自定义数据集加载到库中。

CC-News 数据集包含来自世界各地新闻网站的新闻文章。它包含2017 年 1 月至 2019 年 12 月期间发表的708,241篇英文新闻文章。

下载和准备数据集：

# download and prepare cc_news dataset
dataset = load_dataset("cc_news", split="train")

数据集中只有一个拆分，因此我们需要将其拆分为训练集和测试集：

# split the dataset into training (90%) and testing (10%)
d = dataset.train_test_split(test_size=0.1)
d["train"], d["test"]

你还可以将seed参数传递给该train_test_split()方法，以便在多次运行后它会是相同的集合。

输出：

(Dataset({
     features: ['title', 'text', 'domain', 'date', 'description', 'url', 'image_url'],
     num_rows: 637416
 }), Dataset({
     features: ['title', 'text', 'domain', 'date', 'description', 'url', 'image_url'],
     num_rows: 70825
 }))

让我们看看它的样子：

for t in d["train"]["text"][:3]:
  print(t)
  print("="*50)

输出（剥离）：

Pretty sure women wish men did this better too!!
Q: A recent survey showed that 1/3 of men wish they did THIS better. What is...<STRIPPED>
==================================================
× GoDaddy boots neo-Nazi site after a derogatory story on the Charlottesville victim
The Daily Stormer, a white supremacist and neo-Nazi website,...<STRIPPED>
==================================================
French bank Natixis under investigation over subprime losses
PARIS, Feb 15 Natixis has been placed under formal investigation...<STRIPPED>

如前所述，如果你有自定义数据集，你可以按照上面设置要加载的数据集的链接进行操作，或者LineByLineTextDataset如果你的自定义数据集是一个文本文件，其中所有句子都由一个新的分隔符分隔，则你可以使用该类线。

但是，比使用LineByLineTextDataset设置自定义数据集更好的方法是使用split命令或任何其他 Python 代码将文本文件拆分为多个块文件，然后使用load_dataset()我们上面所做的加载它们，如下所示：

# if you have huge custom dataset separated into files
# load the splitted files
files = ["train1.txt", "train2.txt"] # train3.txt, etc.
dataset = load_dataset("text", data_files=files, split="train")

如果你的自定义数据是一个大文件，那么你应该split在使用该load_dataset()函数加载它们之前将其分成几个文本文件（例如在Linux或Colab上使用该命令），因为如果超过内存，运行时会崩溃.

Transformers训练BERT示例：训练分词器

Python使用Transformers训练BERT？接下来，我们需要训练我们的分词器。为此，我们需要将我们的数据集写入文本文件，因为这就是tokenizers 库要求输入的内容：

# if you want to train the tokenizer from scratch (especially if you have custom
# dataset loaded as datasets object), then run this cell to save it as files
# but if you already have your custom data as text files, there is no point using this
def dataset_to_text(dataset, output_filename="data.txt"):
  """Utility function to save dataset text to disk,
  useful for using the texts to train the tokenizer 
  (as the tokenizer accepts files)"""
  with open(output_filename, "w") as f:
    for t in dataset["text"]:
      print(t, file=f)

# save the training set to train.txt
dataset_to_text(d["train"], "train.txt")
# save the testing set to test.txt
dataset_to_text(d["test"], "test.txt")

上述代码单元格的主要目的是将数据集对象保存为文本文件。如果你已经将数据集作为文本文件，则应跳过此步骤。接下来，让我们定义一些参数：

special_tokens = [
  "[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "<S>", "<T>"
]
# if you want to train the tokenizer on both sets
# files = ["train.txt", "test.txt"]
# training the tokenizer on the training set
files = ["train.txt"]
# 30,522 vocab is BERT's default vocab size, feel free to tweak
vocab_size = 30_522
# maximum sequence length, lowering will result to faster training (when increasing batch size)
max_length = 512
# whether to truncate
truncate_longer_samples = True

复制该files列表是传递给分词器进行训练的文件列表。vocab_size是标记的词汇量。max_length是最大序列长度。现在让我们训练分词器：

# initialize the WordPiece tokenizer
tokenizer = BertWordPieceTokenizer()
# train the tokenizer
tokenizer.train(files=files, vocab_size=vocab_size, special_tokens=special_tokens)
# enable truncation up to the maximum 512 tokens
tokenizer.enable_truncation(max_length=max_length)

由于这是 BERT，因此默认标记器是WordPiece。因此，我们BertWordPieceTokenizer()从tokenizers库中初始化tokenizer 类并使用该train()方法对其进行训练。现在让我们保存它：

model_path = "pretrained-bert"
# make the directory if not already there
if not os.path.isdir(model_path):
  os.mkdir(model_path)
# save the tokenizer  
tokenizer.save_model(model_path)
# dumping some of the tokenizer config to config file, 
# including special tokens, whether to lower case and the maximum sequence length
with open(os.path.join(model_path, "config.json"), "w") as f:
  tokenizer_cfg = {
      "do_lower_case": True,
      "unk_token": "[UNK]",
      "sep_token": "[SEP]",
      "pad_token": "[PAD]",
      "cls_token": "[CLS]",
      "mask_token": "[MASK]",
      "model_max_length": max_length,
      "max_len": max_length,
  }
  json.dump(tokenizer_cfg, f)

该tokenizer.save_model()方法将词汇文件保存到该路径中，我们还手动保存了一些标记器配置，例如特殊标记：

  • unk_token: 代表词表外标记的特殊标记，即使标记器是 WordPiece 标记器，unk标记也不是不可能，而是很少见。
  • sep_token: 在同一个输入中分隔两个不同句子的特殊标记。
  • pad_token: 用于填充未达到最大序列长度的句子的特殊标记（因为标记数组必须具有相同的大小）。
  • cls_token：表示输入类的特殊标记。
  • mask_token：这是我们用于掩码语言建模 (MLM) 预训练任务的掩码标记。

在分词器的训练完成后（应该需要几分钟），让我们现在加载它：

# when the tokenizer is trained and configured, load it as BertTokenizerFast
tokenizer = BertTokenizerFast.from_pretrained(model_path)

标记数据集

Transformers如何训练BERT？现在我们已经准备好了分词器，下面的代码负责对数据集进行分词：

def encode_with_truncation(examples):
  """Mapping function to tokenize the sentences passed with truncation"""
  return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=max_length, return_special_tokens_mask=True)

def encode_without_truncation(examples):
  """Mapping function to tokenize the sentences passed without truncation"""
  return tokenizer(examples["text"], return_special_tokens_mask=True)

# the encode function will depend on the truncate_longer_samples variable
encode = encode_with_truncation if truncate_longer_samples else encode_without_truncation

# tokenizing the train dataset
train_dataset = d["train"].map(encode, batched=True)
# tokenizing the testing dataset
test_dataset = d["test"].map(encode, batched=True)
if truncate_longer_samples:
  # remove other columns and set input_ids and attention_mask as 
  train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])
  test_dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])
else:
  test_dataset.set_format(columns=["input_ids", "attention_mask", "special_tokens_mask"])
  train_dataset.set_format(columns=["input_ids", "attention_mask", "special_tokens_mask"])
train_dataset, test_dataset

encode我们用来标记数据集的回调取决于truncate_longer_samples布尔变量。如果设置为True，那么我们截断超过最大序列长度（max_length参数）的句子。否则，我们不会。

接下来，在设置truncate_longer_samples为的情况下False，我们需要将未截断的样本连接在一起并将它们切成固定大小的向量，因为模型在训练期间需要固定大小的序列：

# Main data processing function that will concatenate all texts from our dataset and generate chunks of
# max_seq_length.
def group_texts(examples):
    # Concatenate all texts.
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
    # customize this part to your needs.
    if total_length >= max_length:
        total_length = (total_length // max_length) * max_length
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + max_length] for i in range(0, total_length, max_length)]
        for k, t in concatenated_examples.items()
    }
    return result
# Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a
# remainder for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value
# might be slower to preprocess.
#
# To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
# https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
if not truncate_longer_samples:
  train_dataset = train_dataset.map(group_texts, batched=True, batch_size=2_000,
                                    desc=f"Grouping texts in chunks of {max_length}")
  test_dataset = test_dataset.map(group_texts, batched=True, batch_size=2_000,
                                  num_proc=4, desc=f"Grouping texts in chunks of {max_length}")

上面的大部分代码是从Huggingface Transformers examples的run_mlm.py脚本中带来的，所以这实际上是由库本身使用的。

如果你不想连接所有文本，然后将它们拆分为 512 个标记的块，请确保设置truncate_longer_samples为True，这样它就会将每一行视为一个单独的样本，而不管其长度如何。请注意，如果你设置truncate_longer_samples为True，则根本不会执行上述代码单元格。

Transformers训练BERT示例：加载模型

在本教程中，我们选择了 BERT，但你可以随意选择 Huggingface 变压器库支持的任何变压器模型，例如RobertaForMaskedLM或DistilBertForMaskedLM：

# initialize the model with the config
model_config = BertConfig(vocab_size=vocab_size, max_position_embeddings=max_length)
model = BertForMaskedLM(config=model_config)

我们使用 初始化模型配置BertConfig，并传递词汇量大小以及最大序列长度。然后我们将配置传递给以BertForMaskedLM初始化模型本身。

预训练

Python使用Transformers训练BERT？在我们开始预训练我们的模型之前，我们需要一种方法来为掩码语言模型 (MLM) 任务随机掩码数据集中的标记。幸运的是，该库通过简单地构造一个DataCollatorForLanguageModeling对象使我们很容易做到这一点：

# initialize the data collator, randomly masking 20% (default is 15%) of the tokens for the Masked Language
# Modeling (MLM) task
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.2
)

我们将tokenizer和 设置mlm为True，并将 设置mlm_probability为0.2以[MASK]20% 的概率随机用令牌替换每个令牌。

接下来，让我们初始化我们的训练参数：

training_args = TrainingArguments(
    output_dir=model_path,          # output directory to where save model checkpoint
    evaluation_strategy="steps",    # evaluate each `logging_steps` steps
    overwrite_output_dir=True,      
    num_train_epochs=10,            # number of training epochs, feel free to tweak
    per_device_train_batch_size=10, # the training batch size, put it as high as your GPU memory fits
    gradient_accumulation_steps=8,  # accumulating the gradients before updating the weights
    per_device_eval_batch_size=64,  # evaluation batch size
    logging_steps=500,             # evaluate, log and save model checkpoints every 1000 step
    save_steps=500,
    # load_best_model_at_end=True,  # whether to load the best model (in terms of loss) at the end of training
    # save_total_limit=3,           # whether you don't have much space so you let only 3 model weights saved in the disk
)

复制每一个参数都在评论中说明，请参阅该TrainingArguments文档的更多细节。现在让我们制作我们的教练：

# initialize the trainer and pass everything to it
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)

我们将训练参数传递给Trainer，以及模型、数据整理器和训练集。我们train()现在只需调用即可开始训练：

# train the model
trainer.train()
[10135/79670 18:53:08 < 129:35:53, 0.15 it/s, Epoch 1.27/10]
Step	Training Loss	Validation Loss
1000	6.904000	6.558231
2000	6.498800	6.401168
3000	6.362600	6.277831
4000	6.251000	6.172856
5000	6.155800	6.071129
6000	6.052800	5.942584
7000	5.834900	5.546123
8000	5.537200	5.248503
9000	5.272700	4.934949
10000	4.915900	4.549236

训练将需要几个小时到几天的时间，具体取决于数据集大小、训练批量大小（即根据你的 GPU 内存大小增加它）和 GPU 速度。

正如你在输出中看到的，模型仍在改进，验证损失仍在下降。一旦验证损失停止减少，你通常必须取消训练。

由于我们已将logging_stepsand设置save_steps为 1000，那么训练器将在每 1000 步后评估并保存模型（即在步骤 x gradient_accumulation_stepx per_device_train_size=  1000x8x10 = 80,000 个样本上进行训练）。结果，我在大约 19 小时的训练，或10000步（即大约1.27 个epochs，或在800,000 个样本上训练）后取消了训练，并开始使用该模型。在下一节中，我们将看到如何使用模型进行推理。

Transformers训练BERT示例：使用模型

Transformers如何训练BERT？在我们使用模型之前，让我们假设当前运行时中没有model和tokenizer变量。因此，我们需要再次加载它们：

# load the model checkpoint
model = BertForMaskedLM.from_pretrained(os.path.join(model_path, "checkpoint-10000"))
# load the tokenizer
tokenizer = BertTokenizerFast.from_pretrained(model_path)

如果你使用的是 Google Colab，那么你必须将检查点保存在 Google Drive 中以备后用，你可以通过设置model_path驱动器路径而不是我们在此处所做的本地路径来实现，只需确保那里有足够的空间.

或者，你可以将你的模型和标记器推送到 Huggingface 集线器中，请查看此有用指南以进行操作。

现在让我们使用我们的模型：

fill_mask = pipeline("fill-mask", model=model, tokenizer=tokenizer)

我们使用简单的管道 API，并同时传递model和tokenizer. 让我们预测一些例子：

# perform predictions
examples = [
  "Today's most trending hashtags on [MASK] is Donald Trump",
  "The [MASK] was cloudy yesterday, but today it's rainy.",
]
for example in examples:
  for prediction in fill_mask(example):
    print(f"{prediction['sequence']}, confidence: {prediction['score']}")
  print("="*50)

输出：

today's most trending hashtags on twitter is donald trump, confidence: 0.1027069091796875
today's most trending hashtags on monday is donald trump, confidence: 0.09271949529647827
today's most trending hashtags on tuesday is donald trump, confidence: 0.08099588006734848
today's most trending hashtags on facebook is donald trump, confidence: 0.04266013577580452
today's most trending hashtags on wednesday is donald trump, confidence: 0.04120611026883125
==================================================
the weather was cloudy yesterday, but today it's rainy., confidence: 0.04445931687951088
the day was cloudy yesterday, but today it's rainy., confidence: 0.037249673157930374
the morning was cloudy yesterday, but today it's rainy., confidence: 0.023775646463036537
the weekend was cloudy yesterday, but today it's rainy., confidence: 0.022554103285074234
the storm was cloudy yesterday, but today it's rainy., confidence: 0.019406016916036606
==================================================

这令人印象深刻，我已经取消了训练，模型仍然产生了有趣的结果！

结论

Python使用Transformers训练BERT？你有一个使用 Huggingface 库预训练 BERT 或其他转换器的完整代码，下面是一些提示：

  • 如上所述，训练速度将取决于 GPU 速度、数据集中的样本数量和批量大小。我已将训练批次大小设置为 10，因为这是它可以在 Colab 上容纳我的 GPU 内存的最大值。如果你有更多内存，请确保增加它，以便显着提高训练速度。
  • 在训练期间，如果你看到验证损失开始增加，请确保记住发生最低验证损失的检查点，以便你可以稍后加载该检查点以供使用。如果你不想跟踪损失，你也可以设置load_best_model_at_end为True，因为它会在训练结束时加载损失方面的最佳权重。
  • 词汇量是根据原始 BERT 配置选择的，因为它的大小为30,522，如果你觉得数据集的语言词汇量很大，可以随意增加它，或者你可以对此进行试验。
  • 如果你设置truncate_longer_samples为False，则代码假定你在一个句子（即行）上有更大的文本，你会注意到处理时间要长得多，尤其是如果你batch_size在map()方法上设置了一个大。如果处理需要很多小时，那么你可以设置truncate_longer_samples为，True以便截断超过max_length标记的句子，或者你可以在使用该save_to_disk()方法处理后保存数据集，以便处理一次并加载多次。

如果你对为下游任务（如文本分类）微调 BERT 感兴趣，那么本教程将指导你完成它。

其他相关教程：

  • 对话式 AI 聊天机器人与 Python 中的 Transformers
  • 如何在 Python 中使用 Transformer 执行文本摘要

在此处查看完整代码。

木子山

发表评论

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen: