1美元从零开始训练Bert，手把手教你优雅地薅谷歌云TPU羊毛

大数据文摘出品

来源：Google Colab

编译：武帅、曹培信

 

2018年10月，Google AI团队推出了Bert，可以说Bert一出生就自带光环。

 

在斯坦福大学机器阅读理解水平测试SQuAD1.1中，Bert在全部两个衡量指标上，全面超越人类表现。并且在另外11种不同NLP测试中均创造了历史以来最好成绩，将GLUE基准提升7.6%，将MultiNLI的准确率提提升5.6%。

 

然而这个拥有12层神经网络的“多头怪”（这里指BERT-Base，BERT-Large有24层），在4个 Cloud TPU 上需要训练 4 天（BERT-Large需要16个Cloud TPU），如此高的训练成本让许多想尝试的同学望而却步。

 

不过，谷歌也给广大程序员带来了福音！我们可以借助谷歌云TPU训练Bert了！并且只需要花费1美元，在Google Colab上还出了完整的教程！

 

快跟文摘菌一起薅谷歌羊毛！

 

在本次实验中，我们将借助谷歌云，在任意文本数据上预训练当下最先进的NLP模型—BERT。

 

BERT模型起源：

https://arxiv.org/abs/1810.04805?source=post_page

 

本指南包含了模型预训练的所有阶段，包括：

• 搭建训练环境

• 下载原始文本数据

• 文本数据预处理

• 学习新词汇表

• 切分预训练数据

• 将数据和模型存储到谷歌云

• 在云TPU上训练模型

 

先回答几个问题

 

这份指南有什么用？

 

借助本指南，你可以在任意文本数据上训练BERT模型。特别是当开源社区没有你需要的语言或示例的预训练模型时，它会帮助到你。

 

谁需要这份指南？

 

这份指南适用于对BERT感兴趣但对当前可用的开源模型的性能并不满意的NLP研究人员。

 

我该如何开始？

 

要想长时间地保存训练数据和模型，你需要一个谷歌云端存储分区（Google Cloud Storage Bucket，GCSB）。请按照这份谷歌云TPU快速入门指南创建一个谷歌云平台账户和谷歌云存储分区。

 

谷歌云TPU快速入门指南：

https://cloud.google.com/tpu/docs/quickstart?source=post_page

 

每一个谷歌云的新用户都可获得300美元的免费金额。

 

链接：

https://cloud.google.com/free/?source=post_page

 

本教程的1到5步出于演示目的，在没有谷歌云存储的情况下也能进行。但是，在这种情况下，你将无法训练模型。

 

需要什么？

 

在第二代TPU（TPUv2）上预训练一个BERT模型大约需要54小时。Google Colab并不是为执行此类需要长时间运行的任务而设计的，它每隔8小时便会中断训练过程。因此，为了训练过程不被中断，你需要使用付费的抢占式的TPUv2。

 

译者注：Google Colab，谷歌免费提供的用于机器学习的平台

相关链接：

https://www.jianshu.com/p/000d2a9d36a0

 

译者注：抢占式，一种进程调度方式，允许将逻辑上可继续运行的进程暂停，适合通用系统。

相关链接：

https://blog.youkuaiyun.com/qq_34173549/article/details/79936219

 

也就是说，在撰写本文时（2019年5月9日），通过Google Colab提供的一块TPU，花费大约1美元，就可以在谷歌云上存储所需要的数据和模型，并预训练一个BERT模型。

 

我该如何遵循指南？

 

下面的代码是Python和Bash的组合。它在Colab Jupyter环境中运行。因此，它可以很方便地在那里运行。

 

然而，除了实际的模型训练部分之外，本指南列出的其他步骤都可以在单独的机器上运行。特别是当你的数据集过大或者十分私密而无法在Colab环境中进行预处理时，这就显得十分有用了。

 

好的，给我看看代码

 

代码链接：

https://colab.research.google.com/drive/1nVn6AFpQSzXBt8_ywfx6XR8ZfQXlKGAz?source=post_page#scrollTo=ODimOhBR05yR

 

我需要修改代码吗？

 

代码中唯一需要你修改的地方就是谷歌云存储的账户名。其他的地方默认就好。

 

还有别的么？

 

说句题外话，除了这个程序，我还发布了一个训练好的俄罗斯语BERT模型。

 

下载链接：

https://storage.googleapis.com/bert_resourses/russian_uncased_L-12_H-768_A-12.zip?source=post_page

 

我希望相关研究人员可以发布其他语言的预训练模型。这样就可以改善我们每个人的NLP环境。现在，让我们进入正题吧！

   

第1步：搭建训练环境

 

首先，我们导入需要用到的包。

 

在Jupyter Notebook中可以通过使用一个感叹号‘！’直接执行bash命令。如下所示：

 

!pip install sentencepiece!git clone https://github.com/google-research/bert

 

整个演示过程我将会用同样的方法使用几个bash命令。

 

现在，让我们导入包并在谷歌云中自行授权。

import osimport sysimport jsonimport nltkimport randomimport loggingimport tensorflow as tfimport sentencepiece as spm

from glob import globfrom google.colab import auth, drivefrom tensorflow.keras.utils import Progbar

sys.path.append("bert")

from bert import modeling, optimization, tokenizationfrom bert.run_pretraining import input_fn_builder, model_fn_builder

auth.authenticate_user()  # configure logginglog = logging.getLogger('tensorflow')log.setLevel(logging.INFO)

# create formatter and add it to the handlersformatter = logging.Formatter('%(asctime)s :  %(message)s')sh = logging.StreamHandler()sh.setLevel(logging.INFO)sh.setFormatter(formatter)log.handlers = [sh]

if 'COLAB_TPU_ADDR' in os.environ:  log.info("Using TPU runtime")  USE_TPU = True  TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']

  with tf.Session(TPU_ADDRESS) as session:    log.info('TPU address is ' + TPU_ADDRESS)    # Upload credentials to TPU.    with open('/content/adc.json', 'r') as f:      auth_info = json.load(f)    tf.contrib.cloud.configure_gcs(session, credentials=auth_info)    else:  log.warning('Not connected to TPU runtime')  USE_TPU = False

Setting up BERT training environment

搭建BERT训练环境

 

第2步：获取数据

 

接下来我们获取文本数据语料库。这次实验我们将采用OpenSubtitles 数据集。

 

链接：

https://www.opensubtitles.org/en/?source=post_page

 

该数据集有65种语言可以使用。

 

链接：

http://opus.nlpl.eu/OpenSubtitles-v2016.php?source=post_page

 

与更常用的文本数据集（如维基百科）不同，该数据集并不需要进行任何复杂的数据预处理。它也预先格式化了，每行一个句子，便于后续处理。

 

你也可以通过设置相应的语言代码来使用该数据集。

AVAILABLE =  {'af','ar','bg','bn','br','bs','ca','cs',              'da','de','el','en','eo','es','et','eu',              'fa','fi','fr','gl','he','hi','hr','hu',              'hy','id','is','it','ja','ka','kk','ko',              'lt','lv','mk','ml','ms','nl','no','pl',              'pt','pt_br','ro','ru','si','sk','sl','sq',              'sr','sv','ta','te','th','tl','tr','uk',              'ur','vi','ze_en','ze_zh','zh','zh_cn',              'zh_en','zh_tw','zh_zh'}

LANG_CODE = "en" #@param {type:"string"}

assert LANG_CODE in AVAILABLE, "Invalid language code selected"

!wget http://opus.nlpl.eu/download.php?f=OpenSubtitles/v2016/mono/OpenSubtitles.raw.'$LANG_CODE'.gz -O dataset.txt.gz!gzip -d dataset.txt.gz!tail dataset.txt

下载OPUS数据

 

出于演示目的，我们默认只使用语料库的一小部分。

 

在实际训练模型时，请务必取消选中DEMO_MODE复选框以使用大100倍的数据集。

 

请放心，一亿行语句足以训练出一个相当不错的BERT模型。

DEMO_MODE = True #@param {type:"boolean"}if DEMO_MODE: CORPUS_SIZE = 1000000else: CORPUS_SIZE = 100000000 #@param {type: "integer"}

!(head -n $CORPUS_SIZE dataset.txt) > subdataset.txt!mv subdataset.txt dataset.txt

 

第3步：文本预处理

 

我们下载的原始文本数据包含了标点符号，大写字母以及非UTF编码的符号，这些都需要提前删除。在模型推断时，我们也需要对新数据集采取同样的做法。

 

如果你的用例需要不同的预处理方法（例如在模型推断时大写字母或者标点符号是需要保留的），那么就修改代码中的函数以满足你的需求。

regex_tokenizer = nltk.RegexpTokenizer("\w+")

def normalize_text(text):  # lowercase text  text = str(text).lower()  # remove non-UTF  text = text.encode("utf-8", "ignore").decode()  # remove punktuation symbols  text = " ".join(regex_tokenizer.tokenize(text))  return text

def count_lines(filename):  count = 0  with open(filename) as fi:    for line in fi:      count += 1  return count

定义预处理例程

 

现在让我们对整个数据集进行预处理吧。

RAW_DATA_FPATH = "dataset.txt" #@param {type: "string"}PRC_DATA_FPATH = "proc_dataset.txt" #@param {type: "string"}

# apply normalization to the dataset# this will take a minute or two

total_lines = count_lines(RAW_DATA_FPATH)bar = Progbar(total_lines)

with open(RAW_DATA_FPATH,encoding="utf-8") as fi:  with open(PRC_DATA_FPATH, "w",encoding="utf-8") as fo:    for l in fi:      fo.write(normalize_text(l)+"\n")      bar.add(1)

 应用预处理

 

第4步：构建词汇表

 

下一步，我们将学习一个新的词汇表，用于表示我们的数据集。

 

因为 BERT 论文中使用了谷歌内部未开源的 WordPiece 分词器，因此，这里我们只能使用一元文法模式（unigram mode）下开源的 SentencePiece 分词器了。

 

链接：

https://github.com/google/sentencepiece?source=post_page

 

虽然它与BERT并不直接兼容，但我们可以通过一个小技巧让它工作。

 

SentencePiece需要相当多的运行内存（RAM），因此在Colab上运行整个数据集会导致内核崩溃。为避免这一情况发生，我们将随机地对数据集的一小部分进行子采样，从而构建词汇表。当然，也可以使用运行内存更大的计算机来执行此步骤。这完全取决于你。

 

此外，SentencePiece默认将BOS和EOS控制符号添加到词汇表中。我们可以通过手动地把它们的词id设为-1来禁用它们。

 

VOC_SIZE的典型值介于32000到128000之间。我们将保留NUM_PLACEHOLDERS标记，以防有人想在训练前的阶段完成后更新词汇和微调模型。

MODEL_PREFIX = "tokenizer" #@param {type: "string"}VOC_SIZE = 32000 #@param {type:"integer"}SUBSAMPLE_SIZE = 12800000 #@param {type:"integer"}NUM_PLACEHOLDERS = 256 #@param {type:"integer"}

SPM_COMMAND = ('--input={} --model_prefix={} '               '--vocab_size={} --input_sentence_size={} '               '--shuffle_input_sentence=true '                '--bos_id=-1 --eos_id=-1').format(               PRC_DATA_FPATH, MODEL_PREFIX,                VOC_SIZE - NUM_PLACEHOLDERS, SUBSAMPLE_SIZE)

spm.SentencePieceTrainer.Train(SPM_COMMAND)

学习SentencePiece词汇表

 

现在，看看我们是如何使SentencePiece在BERT模型中工作的。

 

下面是官方仓库的一个英语BERT预训练模型中通过WordPiece词汇表标记后的语句。

链接：

https://github.com/google-research/bert?source=post_page

 

模型下载：

https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip?source=post_page

>>> wordpiece.tokenize("Colorless geothermal substations are generating furiously")
['color','##less','geo','##thermal','sub','##station','##s','are','generating','furiously']

 

我们看到，WordPiece分词器在两个单词中间以“##”在前的形式预设了子词。单词开头的子词并没有发生变化。如果子词出现在开头和单词的中间，则两个版本（带或不带‘##’）都会添加到词汇表中。

 

SentencePiece创建了两个文件：tokenizer.model and tokenizer.vocab。让我们来看看学到的词汇：

def read_sentencepiece_vocab(filepath):  voc = []  with open(filepath, encoding='utf-8') as fi:    for line in fi:      voc.append(line.split("\t")[0])  # skip the first <unk> token  voc = voc[1:]  return voc

snt_vocab = read_sentencepiece_vocab("{}.vocab".format(MODEL_PREFIX))print("Learnt vocab size: {}".format(len(snt_vocab)))print("Sample tokens: {}".format(random.sample(snt_vocab, 10)))

读取学习后的SentencePiece词汇表

 

给出结果：

Learnt vocab size: 31743Sample tokens: ['▁cafe', '▁slippery', 'xious', '▁resonate', '▁terrier', '▁feat', '▁frequencies', 'ainty', '▁punning', 'modern']

 

我们观察到，SentencePiece和WordPiece给出的结果完全相反。从这篇文档中可以看出SentencePiece首先用元符号“_”(UnicodeMath编码字符：U+2581)替换掉空格，如 “Hello World”被替换为成：

 

Hello▁World.

 

链接：

https://github.com/google/sentencepiece/blob/master/README.md?source=post_page

 

然后，此文本被分割为小块，如下所示：

 

[Hello] [▁Wor] [ld] [.]

 

在空格之后出现的子词（也是大多数单词的开头）通常前面加上了 ‘_’，而其它的并没有变化。这不包括那些仅出现在句子开头而不是其他地方的子词。然而，这些情况很少发生。

 

因此，为了获得类似于WordPiece的词汇表，我们需要进行一个简单的转换，将那些‘_’符号删除并将‘##’符号添加到不含它的标记中。

 

我们还添加了一些BERT架构所需要的特殊控制符号。按照惯例，我们把它们放在了词汇表的开头。

 

此外，我们也在词汇表中添加了一些占位符标记。

 

如果某人希望用新的特定任务的标记来更新模型时，以上那些做法就十分有用了。此时，将原先的占位符标记替换为新的标记，预训练数据就会重新生成，模型也会在新数据上进行微调。

def parse_sentencepiece_token(token):    if token.startswith("▁"):        return token[1:]    else:        return "##" + token        bert_vocab = list(map(parse_sentencepiece_token, snt_vocab))

ctrl_symbols = ["[PAD]","[UNK]","[CLS]","[SEP]","[MASK]"]bert_vocab = ctrl_symbols + bert_vocab

bert_vocab += ["[UNUSED_{}]".format(i) for i in range(VOC_SIZE - len(bert_vocab))]print(len(bert_vocab))

转换词汇表以用于BERT

 

最后，我们将获得的词汇表写入文件。

VOC_FNAME = "vocab.txt" #@param {type:"string"}

with open(VOC_FNAME, "w") as fo:  for token in bert_vocab:    fo.write(token+"\n")

将词汇表写入文件

 

现在，让我们看看新词汇表在实践中是如何运作的：

>>> testcase = "Colorless geothermal substations are generating furiously">>> bert_tokenizer = tokenization.FullTokenizer(VOC_FNAME)>>> bert_tokenizer.tokenize(testcase)['color','##less','geo','##ther','##mal','sub','##station','##s','are','generat','##ing','furious','##ly']

 

第5步：生成预训练数据

 

借助于手头的词汇表，我们已经可以生成BERT模型的预训练数据了。因为我们的数据集可能非常大，所以我们将其切分：

mkdir ./shardssplit -a 4 -l 256000 -d $PRC_DATA_FPATH ./shards/shard_

切分数据集

 

现在，对于每个切片，我们需要从BERT仓库中调用create_pretraining_data.py 脚本。为此，我们使用xargs命令。

 

在我们开始生成数据前，我们需要设置一些参数并传递给脚本。你可以在README文档中找到有关它们含义的更多信息。

 

链接：

https://github.com/google-research/bert/blob/master/README.md?source=post_page

MAX_SEQ_LENGTH = 128 #@param {type:"integer"}MASKED_LM_PROB = 0.15 #@paramMAX_PREDICTIONS = 20 #@param {type:"integer"}DO_LOWER_CASE = True #@param {type:"boolean"}

PRETRAINING_DIR = "pretraining_data" #@param {type:"string"}# controls how many parallel processes xargs can createPROCESSES = 2 #@param {type:"integer"}

定义预训练数据的参数

 

运行这一步可能需要相当长的时间，具体取决于数据集的大小。

XARGS_CMD = ("ls ./shards/ | "             "xargs -n 1 -P {} -I{} "             "python3 bert/create_pretraining_data.py "             "--input_file=./shards/{} "             "--output_file={}/{}.tfrecord "             "--vocab_file={} "             "--do_lower_case={} "             "--max_predictions_per_seq={} "             "--max_seq_length={} "             "--masked_lm_prob={} "             "--random_seed=34 "             "--dupe_factor=5")

XARGS_CMD = XARGS_CMD.format(PROCESSES, '{}', '{}', PRETRAINING_DIR, '{}',                              VOC_FNAME, DO_LOWER_CASE,                              MAX_PREDICTIONS, MAX_SEQ_LENGTH, MASKED_LM_PROB)                             tf.gfile.MkDir(PRETRAINING_DIR)!$XARGS_CMD

创建预训练数据

 

第6步：创建持久化存储

 

为了保存我们来之不易的财产，我们将其保存在谷歌云存储上。如果你已经创建了谷歌云存储分区，那么这是很容易实现的。

 

我们将在谷歌云存储上设置两个目录：一个用于数据，一个用于模型。在模型目录下，我们将放置模型词汇表和配置文件。

 

在继续操作之前，请在此配置你的BUCKET_NAME变量，否则你将无法训练模型。

BUCKET_NAME = "bert_resourses" #@param {type:"string"}MODEL_DIR = "bert_model" #@param {type:"string"}tf.gfile.MkDir(MODEL_DIR)

if not BUCKET_NAME:  log.warning("WARNING: BUCKET_NAME is not set. "              "You will not be able to train the model.")

配置GCS bucket名称

 

下面是BERT的超参数配置示例。若更改将自担风险！

# use this for BERT-base

bert_base_config = {  "attention_probs_dropout_prob": 0.1,   "directionality": "bidi",   "hidden_act": "gelu",   "hidden_dropout_prob": 0.1,   "hidden_size": 768,   "initializer_range": 0.02,   "intermediate_size": 3072,   "max_position_embeddings": 512,   "num_attention_heads": 12,   "num_hidden_layers": 12,   "pooler_fc_size": 768,   "pooler_num_attention_heads": 12,   "pooler_num_fc_layers": 3,   "pooler_size_per_head": 128,   "pooler_type": "first_token_transform",   "type_vocab_size": 2,   "vocab_size": VOC_SIZE}

with open("{}/bert_config.json".format(MODEL_DIR), "w") as fo:  json.dump(bert_base_config, fo, indent=2)  with open("{}/{}".format(MODEL_DIR, VOC_FNAME), "w") as fo:  for token in bert_vocab:    fo.write(token+"\n")

配置BERT的超参数并保存到磁盘

 

现在，我们将我们的成果存放到谷歌云存储。

 

if BUCKET_NAME:  !gsutil -m cp -r $MODEL_DIR $PRETRAINING_DIR gs://$BUCKET_NAME

上传成果到谷歌云存储

 

第7步：训练模型

 

 我们已准备好开始训练我们的模型了。

 

建议过去步骤中的一些参数不要更改，以便于快速重启训练过程。

 

确保整个实验所设置的参数完全相同。

BUCKET_NAME = "bert_resourses" #@param {type:"string"}MODEL_DIR = "bert_model" #@param {type:"string"}PRETRAINING_DIR = "pretraining_data" #@param {type:"string"}VOC_FNAME = "vocab.txt" #@param {type:"string"}

# Input data pipeline configTRAIN_BATCH_SIZE = 128 #@param {type:"integer"}MAX_PREDICTIONS = 20 #@param {type:"integer"}MAX_SEQ_LENGTH = 128 #@param {type:"integer"}MASKED_LM_PROB = 0.15 #@param

# Training procedure configEVAL_BATCH_SIZE = 64LEARNING_RATE = 2e-5TRAIN_STEPS = 1000000 #@param {type:"integer"}SAVE_CHECKPOINTS_STEPS = 2500 #@param {type:"integer"}NUM_TPU_CORES = 8

if BUCKET_NAME:  BUCKET_PATH = "gs://{}".format(BUCKET_NAME)else:  BUCKET_PATH = "."

BERT_GCS_DIR = "{}/{}".format(BUCKET_PATH, MODEL_DIR)DATA_GCS_DIR = "{}/{}".format(BUCKET_PATH, PRETRAINING_DIR)

VOCAB_FILE = os.path.join(BERT_GCS_DIR, VOC_FNAME)CONFIG_FILE = os.path.join(BERT_GCS_DIR, "bert_config.json")

INIT_CHECKPOINT = tf.train.latest_checkpoint(BERT_GCS_DIR)

bert_config = modeling.BertConfig.from_json_file(CONFIG_FILE)input_files = tf.gfile.Glob(os.path.join(DATA_GCS_DIR,'*tfrecord'))

log.info("Using checkpoint: {}".format(INIT_CHECKPOINT))log.info("Using {} data shards".format(len(input_files)))

配置训练过程

 

准备训练运行配置，建立评估器和输入函数，启动BERT。

model_fn = model_fn_builder(      bert_config=bert_config,      init_checkpoint=INIT_CHECKPOINT,      learning_rate=LEARNING_RATE,      num_train_steps=TRAIN_STEPS,      num_warmup_steps=10,      use_tpu=USE_TPU,      use_one_hot_embeddings=True)

tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)

run_config = tf.contrib.tpu.RunConfig(    cluster=tpu_cluster_resolver,    model_dir=BERT_GCS_DIR,    save_checkpoints_steps=SAVE_CHECKPOINTS_STEPS,    tpu_config=tf.contrib.tpu.TPUConfig(        iterations_per_loop=SAVE_CHECKPOINTS_STEPS,        num_shards=NUM_TPU_CORES,        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2))

estimator = tf.contrib.tpu.TPUEstimator(    use_tpu=USE_TPU,    model_fn=model_fn,    config=run_config,    train_batch_size=TRAIN_BATCH_SIZE,    eval_batch_size=EVAL_BATCH_SIZE)  train_input_fn = input_fn_builder(        input_files=input_files,        max_seq_length=MAX_SEQ_LENGTH,        max_predictions_per_seq=MAX_PREDICTIONS,        is_training=True)

建立评估模型和输入函数

 

启动！

 

estimator.train(input_fn=train_input_fn, max_steps=TRAIN_STEPS)

启动BERT

 

使用默认参数训练模型100万步大约需要53个小时。如果内核出于某种原因重新启动，你可以从最新的检查点继续训练。

 

以上就是在云TPU上从头开始预训练BERT模型的指南。

 

还能做点什么？

 

我们已经训练好了模型，那接下来呢？

 

这是一个全新的话题。你可以做如下几件事：

• 将预训练模型作为通用NLU模块

• 针对某些特定的分类任务微调模型

• 使用BERT作为构件块创建另一个深度学习模型

 

链接：

https://towardsdatascience.com/building-a-search-engine-with-bert-and-tensorflow-c6fdc0186c8a

 

真正有趣的东西还在后面，所以睁大你的眼睛。同时，看看bear-as-a-service这个很棒的项目，将你刚训好的模型部署到线上去吧！

 

项目链接：

https://github.com/hanxiao/bert-as-service?source=post_page

 

最最重要的，不断学习！