BERT学习笔记：run_classifier.py

BERT 源码初探之 run_classifier.py

本文源码来源于 Github上的BERT 项目中的 run_classifier.py 文件。阅读本文需要对Attention Is All You Need以及BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding两篇论文有所了解，以及部分关于深度学习、自然语言处理和Tensorflow的储备知识。

0 前言

• 关于Tensorflow：本文基于谷歌官方在GitHub上公布的BERT预训练模型，基于Tensorflow 1.13.1 运行。有关Tensorflow的部分建议参照官方网站。
• 关于Transformer：Transformer是Google提出的一种完全基于注意力机制的模型，想要对齐进行了解请参照官方论文Attention Is All You Need或者我的另一篇博客Transformer 学习笔记。
• 关于BERT：BERT也是Google提出的一个基于Transformer的预训练网络模型，更多和该模型有关的内容请参照官方论文BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding、官方代码实现Github上的BERT以及我的另一篇博客BERT 学习笔记。

1 简介

2 classifier使用方法

2.1 参数设置

2.1.1 必须参数

源码中提到的必须要设置的参数如下所示：

flags.DEFINE_string(
    "data_dir", None,
    "The input data dir. Should contain the .tsv files (or other data files) "
    "for the task.")

• 名称：data_dir
• 默认值: None
• 注释：输入数据的目录，需要包含任务所需的 .csv 格式（或者其他类型）的文件。

flags.DEFINE_string(
    "bert_config_file", None,
    "The config json file corresponding to the pre-trained BERT model. "
    "This specifies the model architecture.")

• 名称：bert_config_file
• 默认值：None
• 注释：和预训练过的bert模型对应的json格式的配置文件，这指定了模型的架构。

flags.DEFINE_string("task_name", None, "The name of the task to train.")

• 名称：task_name
• 默认值：None
• 注释：要训练的任务名称。

flags.DEFINE_string("vocab_file", None,
                    "The vocabulary file that the BERT model was trained on.")

• 名称：vocab_file
• 默认值：None
• 注释：BERT模型训练所用的词典文件

flags.DEFINE_string(
    "output_dir", None,
    "The output directory where the model checkpoints will be written.")

• 名称：output_dir
• 默认值：None
• 注释：模型checkpoint会被保存到的路径位置。

2.1.2 其他参数

flags.DEFINE_string(
    "init_checkpoint", None,
    "Initial checkpoint (usually from a pre-trained BERT model).")

• 名称：init_checkpoint
• 默认值：None
• 注释：初始的checkpoint（通常来源于一个预训练过的BERT模型）。

flags.DEFINE_bool(
    "do_lower_case", True,
    "Whether to lower case the input text. Should be True for uncased "
    "models and False for cased models.")

• 名称：do_lower_case
• 默认值：True
• 注释：是否把输入的文本全部小写。对于所有的 uncased models（不保留大小写以及重音标记）应该设为True，而对于 cased models（保留大小写以及重音标记，如果你认为这对于你的训练任务是有益的）则应该设为False。

flags.DEFINE_integer(
    "max_seq_length", 128,
    "The maximum total input sequence length after WordPiece tokenization. "
    "Sequences longer than this will be truncated, and sequences shorter "
    "than this will be padded.")

• 名称：max_seq_length
• 默认值：128
• 注释：WordPiece tokenization处理之后的最大输入序列长度。比这个长度长的输入序列将被阶段，而比这个短的输入序列将被补齐。

flags.DEFINE_bool("do_train", False, "Whether to run training.")

• 名称：do_train
• 默认值：False
• 注释：是否运行训练

flags.DEFINE_bool("do_eval", False, "Whether to run eval on the dev set.")

• 名称：do_eval
• 默认值：False
• 注释：是否在验证集上进行计算（eval）。

flags.DEFINE_bool(
    "do_predict", False,
    "Whether to run the model in inference mode on the test set.")

• 名称：do_predict
• 默认值：False
• 注释：是否在测试集上用推断模式来运行模型。

flags.DEFINE_integer("train_batch_size", 32, "Total batch size for training.")

• 名称：train_batch_size
• 默认值：32
• 注释：训练过程中的批尺寸。

flags.DEFINE_integer("eval_batch_size", 8, "Total batch size for eval.")

• 名称：eval_batch_size
• 默认值：8
• 注释：验证过程中的批尺寸。

flags.DEFINE_integer("predict_batch_size", 8, "Total batch size for predict.")

• 名称：predict_batch_size
• 默认值：8
• 注释：预测过程中的批尺寸。

flags.DEFINE_float("learning_rate", 5e-5, "The initial learning rate for Adam.")

• 名称：learning_rate
• 默认值：5e-5
• 注释：Adam算法的初始学习率。

flags.DEFINE_float("num_train_epochs", 3.0,
                   "Total number of training epochs to perform.")

• 名称：num_train_epochs
• 默认值：3.0
• 注释：完成训练所需要遍历训练数据集的次数

flags.DEFINE_float(
    "warmup_proportion", 0.1,
    "Proportion of training to perform linear learning rate warmup for. "
    "E.g., 0.1 = 10% of training.")

• 名称：warmup_proportion
• 默认值：0.1
• 注释：用于warm up的训练步数比例。其中0.1表示训练步数的10%。

flags.DEFINE_integer("save_checkpoints_steps", 1000,
                     "How often to save the model checkpoint.")

• 名称：save_checkpoints_steps
• 默认值：1000
• 注释：多久保存一次模型

flags.DEFINE_integer("iterations_per_loop", 1000,
                     "How many steps to make in each estimator call.")

• 名称：iterations_per_loop
• 默认值：1000
• 注释：每隔多少步来调用一次估计函数。

2.1.3 TPU相关参数

flags.DEFINE_bool("use_tpu", False, "Whether to use TPU or GPU/CPU.")

• 名称：use_tpu
• 默认值：False
• 注释：是否使用TPU

tf.flags.DEFINE_string(
    "tpu_name", None,
    "The Cloud TPU to use for training. This should be either the name "
    "used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 "
    "url.")

• 名称：tpu_name
• 默认值：None
• 注释：用于训练的云TPU，应该是创建云TPU时使用的名称，或者是类似 grpc://ip.address.of.tpu:8470 这种类型的URL。

tf.flags.DEFINE_string(
    "tpu_zone", None,
    "[Optional] GCE zone where the Cloud TPU is located in. If not "
    "specified, we will attempt to automatically detect the GCE project from "
    "metadata.")

• 名称：tpu_zone
• 默认值：None
• 注释：【可选】 云TPU所在的GCE区域，如果没有指定该参数，程序将会自动在元数据中检测。

tf.flags.DEFINE_string(
    "gcp_project", None,
    "[Optional] Project name for the Cloud TPU-enabled project. If not "
    "specified, we will attempt to automatically detect the GCE project from "
    "metadata.")

• 名称：gcp_project
• 默认值：None
• 注释：【可选】 云TPU项目的项目名称，如果不指定程序将自动从元数据中检测。

tf.flags.DEFINE_string("master", None, "[Optional] TensorFlow master URL.")

• 名称：master
• 默认值：None
• 注释：【可选】 Tensorflow主节点的URL。

flags.DEFINE_integer(
    "num_tpu_cores", 8,
    "Only used if `use_tpu` is True. Total number of TPU cores to use.")

• 名称：num_tpu_cores
• 默认值：8
• 注释：只有在 “use_tpu” 参数被设为 True 时才会使用。其含义为使用的TPU核心的数目。

2.2 输入数据相关

2.2.1 输入样例（InputExample）

class InputExample(object):

  def __init__(self, guid, text_a, text_b=None, label=None):
    self.guid = guid
    self.text_a = text_a
    self.text_b = text_b
    self.label = label

该类构造了一个简单的训练/测试的输入序列类型。
参数含义：

• guid：该条样例的唯一ID
• text_a：字符串。第一个序列的未标记化(untokenized)的文本。针对单一序列的训练任务，只有这个序列必须被指定。
• text_b：【可选】 字符串。第二个序列未标记化(untokenized)的文本。 只有在序列对任务中才需要被指定。
• label：【可选】 字符串。该例子的标签。应在训练集和验证集中提供，但是在测试集中不需要被提供。

2.2.2 用于填充的样例(PaddingInputExample)

class PaddingInputExample(object):
    # 这里什么也没有

用来填充训练数据以便让训练数据的大小是batch_size的整数倍。为了在TPU上运行eval或者predict，我们需要把输入数据填充为batch_size的整数倍，因为TPU需要一个固定的batch_size，另一种方法是删除最后一批大小小于batch_size的输入数据，这种做法很糟糕因为这并不会输出所有的生成结果。我们采用这个类而不是None来进行填充因为None会产生没有提示的错误。

2.2.3 输入特征(InputFeatures)

class InputFeatures(object):

  def __init__(self,
               input_ids,
               input_mask,
               segment_ids,
               label_id,
               is_real_example=True):
    self.input_ids = input_ids
    self.input_mask = input_mask
    self.segment_ids = segment_ids
    self.label_id = label_id
    self.is_real_example = is_real_example

数据的一组特征。用于参加训练的数据格式。

• input_ids 是将词映射成id的向量
• input_mask 是论文中提到的mask向量
• segment_ids 用于句子对的标记，单句全为0，句子对则第一句为0第二句为1。
• label_id 是输入数据的标签
• is_real_example 标记此数据是否为真实的输入数据

2.2.4 数据处理器基类（DataProcessor）

class DataProcessor(object):
  """用于序列分类任务的数据转换器的基类"""

  def get_train_examples(self, data_dir):
    """定义了获取训练数据的 InputExample 格式的集合的方法。"""
    raise NotImplementedError()

  def get_dev_examples(self, data_dir):
    """定义了获取验证数据的 InputExample 格式的集合的方法。"""
    raise NotImplementedError()

  def get_test_examples(self, data_dir):
    """定义了获取预测数据的 InputExample 格式的集合的方法。"""
    raise NotImplementedError()

  def get_labels(self):
    """获取这个数据集的标签列表"""
    raise NotImplementedError()

  @classmethod
  def _read_tsv(cls, input_file, quotechar=None):
    """读取用tab分割的数据文件"""
    with tf.gfile.Open(input_file, "r") as f:
      reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
      lines = []
      for line in reader:
        lines.append(line)
      return lines

这个类是一个把数据从文件转换为InputExample格式的数据转换器的基类，针对不同的数据文件我们需要自定义一个类继承自该类来实现定义好的方法。

2.2.4.1 XNLI数据集的数据转换器(XnliProcessor)

class XnliProcessor(DataProcessor):

该类继承自DataProcessor。

  def __init__(self):
    self.language = "zh"

该类的初始化方法，指定了language。

  def get_train_examples(self, data_dir):
    lines = self._read_tsv(
        os.path.join(data_dir, "multinli",
                     "multinli.train.%s.tsv" % self.language))
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "train-%d" % (i)
      text_a = tokenization.convert_to_unicode(line[0])
      text_b = tokenization.convert_to_unicode(line[1])
      label = tokenization.convert_to_unicode(line[2])
      if label == tokenization.convert_to_unicode("contradictory"):
        label = tokenization.convert_to_unicode("contradiction")
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

get_train_examples 方法的实现。tokenization.convert_to_unicode 方法把所有的输入数据转换成 unicode 格式。通过对数据文件中按行读取，把每行的数据组合成 InputExample 格式的输入数据。然后返回整个example的List集合。

  def get_dev_examples(self, data_dir):
    lines = self._read_tsv(os.path.join(data_dir, "xnli.dev.tsv"))
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "dev-%d" % (i)
      language = tokenization.convert_to_unicode(line[0])
      if language != tokenization.convert_to_unicode(self.language):
        continue
      text_a = tokenization.convert_to_unicode(line[6])
      text_b = tokenization.convert_to_unicode(line[7])
      label = tokenization.convert_to_unicode(line[1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

获取验证集方法的实现，基本逻辑和训练集相同。

  def get_labels(self):
    return ["contradiction", "entailment", "neutral"]

获取标签的方法实现。在当前数据集中只有三种标签。

2.2.4.2 MultiNLI数据集的数据转换器(MnliProcessor)

class MnliProcessor(DataProcessor):

该类继承自DataProcessor。

  def get_train_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev_matched.tsv")),
        "dev_matched")

  def get_test_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test_matched.tsv")), "test")

获取三种输入数据的方法都被整合到了_create_examples方法当中。

  def get_labels(self):
    return ["contradiction", "entailment", "neutral"]

获取标签方法的实现。

  def _create_examples(self, lines, set_type):
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, tokenization.convert_to_unicode(line[0]))
      text_a = tokenization.convert_to_unicode(line[8])
      text_b = tokenization.convert_to_unicode(line[9])
      if set_type == "test":
        label = "contradiction"
      else:
        label = tokenization.convert_to_unicode(line[-1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

总体上和XNLI数据集类似，因为把训练/验证/测试集方法整合到了一起，因此用一个输入参数 set_type 来区分要返回的集合类型。

2.2.4.3 MRPC数据集的数据转换器(MrpcProcessor)

class MrpcProcessor(DataProcessor):

  def get_train_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

  def get_test_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    return ["0", "1"]

  def _create_examples(self, lines, set_type):
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
      text_a = tokenization.convert_to_unicode(line[3])
      text_b = tokenization.convert_to_unicode(line[4])
      if set_type == "test":
        label = "0"
      else:
        label = tokenization.convert_to_unicode(line[0])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

内容和MultiNLI数据集转换器高度类似，因此不在此过多解释。

2.2.4.4 CoLA数据集数据转换器(ColaProcessor)

class ColaProcessor(DataProcessor):

  def get_train_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

  def get_test_examples(self, data_dir):
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    return ["0", "1"]

  def _create_examples(self, lines, set_type):
    examples = []
    for (i, line) in enumerate(lines):
      # Only the test set has a header
      if set_type == "test" and i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
      if set_type == "test":
        text_a = tokenization.convert_to_unicode(line[1])
        label = "0"
      else:
        text_a = tokenization.convert_to_unicode(line[3])
        label = tokenization.convert_to_unicode(line[1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
    return examples

内容和MultiNLI数据集高度类似，因此不在此过多解释。

通过对以上四个数据集的转换器源码阅读，我们对BERT的数据转换器会有一个比较透彻的理解，可以借此针对自己的输入数据集来编写对应的数据转换器。

2.2.5 将单一的InputExample转换成对应的InputFeatures格式(convert_single_example方法)

def convert_single_example(ex_index, example, label_list, max_seq_length,
                           tokenizer):

函数定义与输入参数。

  if isinstance(example, PaddingInputExample):
    return InputFeatures(
        input_ids=[0] * max_seq_length,
        input_mask=[0] * max_seq_length,
        segment_ids=[0] * max_seq_length,
        label_id=0,
        is_real_example=False)

如果输入数据是填充样例类，那么就返回一个用0填充的假输入样例。

  label_map = {}
  for (i, label) in enumerate(label_list):
    label_map[label] = i

对输入数据的标签列表进行map映射化。

  tokens_a = tokenizer.tokenize(example.text_a)
  tokens_b = None
  if example.text_b:
    tokens_b = tokenizer.tokenize(example.text_b)

  if tokens_b:
    _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)
  else:
    if len(tokens_a) > max_seq_length - 2:
      tokens_a = tokens_a[0:(max_seq_length - 2)]

对tokens_a和tokens_b进行预处理，首先将其标记化，然后如果tokens_b存在，那么tokens_a和tokens_b的长度就不能超过max_seq_length-3，因为需要加入cls,sep,seq三个符号。如果只有tokens_a，那么tokens_a的长度不能超过 max_seq_length -2 (加入 cls 和 sep符号)。

def _truncate_seq_pair(tokens_a, tokens_b, max_length):
  while True:
    total_length = len(tokens_a) + len(tokens_b)
    if total_length <= max_length:
      break
    if len(tokens_a) > len(tokens_b):
      tokens_a.pop()
    else:
      tokens_b.pop()

通过该方法我们可以把输入的句子对截断成最大长度。

  tokens = []
  segment_ids = []
  tokens.append("[CLS]")
  segment_ids.append(0)
  for token in tokens_a:
    tokens.append(token)
    segment_ids.append(0)
  tokens.append("[SEP]")
  segment_ids.append(0)

  if tokens_b:
    for token in tokens_b:
      tokens.append(token)
      segment_ids.append(1)
    tokens.append("[SEP]")
    segment_ids.append(1)

    input_ids = tokenizer.convert_tokens_to_ids(tokens)

    input_mask = [1] * len(input_ids)

按照论文中的格式对句子进行格式化，并根据词典把输入的单词转换成id。对于真实的token其mask_id是1，填充的是0，我们只关注真实的token。

  while len(input_ids) < max_seq_length:
    input_ids.append(0)
    input_mask.append(0)
    segment_ids.append(0)

  assert len(input_ids) == max_seq_length
  assert len(input_mask) == max_seq_length
  assert len(segment_ids) == max_seq_length

对于长度不够的标签进行补齐，同时确保三个特征向量的长度符合规定。

  label_id = label_map[example.label]
  if ex_index < 5:
    tf.logging.info("*** Example ***")
    tf.logging.info("guid: %s" % (example.guid))
    tf.logging.info("tokens: %s" % " ".join(
        [tokenization.printable_text(x) for x in tokens]))
    tf.logging.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
    tf.logging.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
    tf.logging.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
    tf.logging.info("label: %s (id = %d)" % (example.label, label_id))

  feature = InputFeatures(
      input_ids=input_ids,
      input_mask=input_mask,
      segment_ids=segment_ids,
      label_id=label_id,
      is_real_example=True)
  return feature

根据map获取标签的id,打印出前五个输入数据的信息，然后把输入数据包装成InputFeatures格式并返回。

2.2.6 处理完毕的输入特征数据读写相关

2.2.6.1 把一个InputFeature数据集写入TFRecord文件中(file_based_convert_examples_to_features)

def file_based_convert_examples_to_features(
    examples, label_list, max_seq_length, tokenizer, output_file):

  writer = tf.python_io.TFRecordWriter(output_file)

  for (ex_index, example) in enumerate(examples):
    if ex_index % 10000 == 0:
      tf.logging.info("Writing example %d of %d" % (ex_index, len(examples)))

    feature = convert_single_example(ex_index, example, label_list,
                                     max_seq_length, tokenizer)

    def create_int_feature(values):
      f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
      return f

    features = collections.OrderedDict()
    features["input_ids"] = create_int_feature(feature.input_ids)
    features["input_mask"] = create_int_feature(feature.input_mask)
    features["segment_ids"] = create_int_feature(feature.segment_ids)
    features["label_ids"] = create_int_feature([feature.label_id])
    features["is_real_example"] = create_int_feature(
        [int(feature.is_real_example)])

    tf_example = tf.train.Example(features=tf.train.Features(feature=features))
    writer.write(tf_example.SerializeToString())
  writer.close()

当我们对模型进行训练的时候，一般的流程是把训练数据全部加载到内存当中，然后随着程序的运行来进行训练，这对于小规模的数据没有问题，但是当我们针对大数据集进行训练的时候，就会发现我们的内存并不足以让我们把训练数据全部加载到内存当中，据此tensorflow提供给了开发者一种特殊的二进制格式的数据文件 .record ，和TFRecord有关的内容在此不在赘述。

上面代码的所用就是把输入example转换成inputfeatures格式然后写入record格式的二进制文件，以便训练的时候模型加载数据。

2.2.6.2 创建一个 input_fn 闭包方法传递给TPUEstimator(file_based_input_fn_builder)

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([], tf.int64),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

定义输入特征的名称。

  def _decode_record(record, name_to_features):
    example = tf.parse_single_example(record, name_to_features)

    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

首先把一条数据根据输入特征的名称解码成tf.example，因为tf.example只支持tf.int64，但是TPU只支持tf.int32，所以需要把所有的int64转换成int32。

  def input_fn(params):
    batch_size = params["batch_size"]

    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

真正的输入方法，对于训练数据，我们需要大量的并行读写和打乱顺序，而对于验证数据，我们不希望打乱数据，是否并行也不关心。

def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder):
  
  ......

  return input_fn

执行该函数会返回input_fn方法。

2.3 创建模型

2.3.1 建立一个BERT分类模型(create_model)

def create_model(bert_config, is_training, input_ids, input_mask, segment_ids,
                 labels, num_labels, use_one_hot_embeddings):
  model = modeling.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      token_type_ids=segment_ids,
      use_one_hot_embeddings=use_one_hot_embeddings)

首先，根据modeling.BertModel所需的参数来定义一个分类模型

  output_layer = model.get_pooled_output()

  hidden_size = output_layer.shape[-1].value

  output_weights = tf.get_variable(
      "output_weights", [num_labels, hidden_size],
      initializer=tf.truncated_normal_initializer(stddev=0.02))

  output_bias = tf.get_variable(
      "output_bias", [num_labels], initializer=tf.zeros_initializer())

• 在本次实验中，我们只需要一个整体上的输出，因此调用了 get_pooled_output 方法，如果需要获取字符级别的输出，可以调用 get_sequence_output 方法。
• 通过modeling中的源码注释我们可以知道， sequence_output 的 shape = [batch_size, seq_length, hidden_size] ，而 pooler 就是把 sequence_output 的 shape 转换成 [batch_size, hidden_size] 。因此，隐藏层的大小就是 output_layer的shape属性的最后一个值。
• 然后通过 num_label 以及 hidden_size 来构建输出的权重（weight）和偏置值（bias）。

  with tf.variable_scope("loss"):
    if is_training:
      output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)

    logits = tf.matmul(output_layer, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    probabilities = tf.nn.softmax(logits, axis=-1)
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)

    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
    loss = tf.reduce_mean(per_example_loss)

    return (loss, per_example_loss, logits, probabilities)

• 指定变量的作用域 loss
• 如果在训练中那么就把 output_layer加上 dropout 来防止过拟合。
• logits 的值就是简单的线性函数 y=wx+b
• 预测的结果（probabilities）就是在 logits 上进行 softmax。
• 通过计算交叉着来获取损失函数
• 返回损失函数、每隔样本的差距、线性函数的结果、预测概率

2.3.2 自定义模型估计器(model_fn_builder)

def model_fn_builder(bert_config, num_labels, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings):

  def model_fn(features, labels, mode, params):

  ...

  return model_fn

方法头，该方法返回一个 TPUEstimator 默认调用的 model_fn 方法。

    tf.logging.info("*** Features ***")
    for name in sorted(features.keys()):
      tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

打印特征的名称和shape

    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    label_ids = features["label_ids"]
    is_real_example = None
    if "is_real_example" in features:
      is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32)
    else:
      is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)

根据创建模型的参数来初始化所需的参数

    (total_loss, per_example_loss, logits, probabilities) = create_model(
        bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
        num_labels, use_one_hot_embeddings)

调用创建模型的方法，获得返回的模型

    tvars = tf.trainable_variables()
    initialized_variable_names = {}
    scaffold_fn = None
    if init_checkpoint:
      (assignment_map, initialized_variable_names
      ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
      if use_tpu:

        def tpu_scaffold():
          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
          return tf.train.Scaffold()

        scaffold_fn = tpu_scaffold
      else:
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

• 获得需要训练的变量列表
• 如果存在待初始化的checkpoint那么就计算当前变量和检查点变量的并集
• 根据是否使用TPU来选择从检查点初始化的方法

    tf.logging.info("**** Trainable Variables ****")
    for var in tvars:
      init_string = ""
      if var.name in initialized_variable_names:
        init_string = ", *INIT_FROM_CKPT*"
      tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                      init_string)

打印所有可以训练的变量信息

    output_spec = None
    if mode == tf.estimator.ModeKeys.TRAIN:

      train_op = optimization.create_optimizer(
          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)

      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op,
          scaffold_fn=scaffold_fn)
    elif mode == tf.estimator.ModeKeys.EVAL:

      def metric_fn(per_example_loss, label_ids, logits, is_real_example):
        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
        accuracy = tf.metrics.accuracy(
            labels=label_ids, predictions=predictions, weights=is_real_example)
        loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)
        return {
            "eval_accuracy": accuracy,
            "eval_loss": loss,
        }

      eval_metrics = (metric_fn,
                      [per_example_loss, label_ids, logits, is_real_example])
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          eval_metrics=eval_metrics,
          scaffold_fn=scaffold_fn)
    else:
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          predictions={"probabilities": probabilities},
          scaffold_fn=scaffold_fn)
    return output_spec

根据当前训练状态（训练、验证还是测试）来创建不同的 TPUEstimatorSpec 。

2.3.3 两个附带的函数

input_fn_builder(features, seq_length, is_training, drop_remainder)

convert_examples_to_features(examples, label_list, max_seq_length,
                                 tokenizer)

根据源文件注释显示，这两个方法不在本文件中使用，但是在Colab和依赖他的人使用，因此在此不多加解释。

2.4 main(_) 函数

2.4.1 模型初始化部分

  tf.logging.set_verbosity(tf.logging.INFO)

  processors = {
      "cola": ColaProcessor,
      "mnli": MnliProcessor,
      "mrpc": MrpcProcessor,
      "xnli": XnliProcessor,
  }

设置日志记录等级，数据处理器字典

  tokenization.validate_case_matches_checkpoint(FLAGS.do_lower_case,
                                                FLAGS.init_checkpoint)

  if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
    raise ValueError(
        "At least one of `do_train`, `do_eval` or `do_predict' must be True.")

  bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)

  if FLAGS.max_seq_length > bert_config.max_position_embeddings:
    raise ValueError(
        "Cannot use sequence length %d because the BERT model "
        "was only trained up to sequence length %d" %
        (FLAGS.max_seq_length, bert_config.max_position_embeddings))

• 检查当前的配置和checkpoint的配置是否匹配
• 确保 do_train do_eval do_predict 三个参数至少有一个置为true
• 确保最大句子长度不超过 BERT 支持的最大句子长度

tf.gfile.MakeDirs(FLAGS.output_dir)

根据输出目录参数创建目录

  task_name = FLAGS.task_name.lower()

  if task_name not in processors:
    raise ValueError("Task not found: %s" % (task_name))

  processor = processors[task_name]()

  label_list = processor.get_labels()

  tokenizer = tokenization.FullTokenizer(
      vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)

• 根据任务名称来判断当前任务是否可以执行
• 获取任务对应的数据处理器和标签列表
• 根据词典文件初始化分词器

  tpu_cluster_resolver = None
  if FLAGS.use_tpu and FLAGS.tpu_name:
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
        FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)

  is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2
  run_config = tf.contrib.tpu.RunConfig(
      cluster=tpu_cluster_resolver,
      master=FLAGS.master,
      model_dir=FLAGS.output_dir,
      save_checkpoints_steps=FLAGS.save_checkpoints_steps,
      tpu_config=tf.contrib.tpu.TPUConfig(
          iterations_per_loop=FLAGS.iterations_per_loop,
          num_shards=FLAGS.num_tpu_cores,
          per_host_input_for_training=is_per_host))

如果使用 TPU 那么就初始化和 TPU 相关的内容

  train_examples = None
  num_train_steps = None
  num_warmup_steps = None
  if FLAGS.do_train:
    train_examples = processor.get_train_examples(FLAGS.data_dir)
    num_train_steps = int(
        len(train_examples) / FLAGS.train_batch_size * FLAGS.num_train_epochs)
    num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)

如果要进行训练，那么就获取训练数据来计算一共需要的训练步数以及 warmup 的步数

  model_fn = model_fn_builder(
      bert_config=bert_config,
      num_labels=len(label_list),
      init_checkpoint=FLAGS.init_checkpoint,
      learning_rate=FLAGS.learning_rate,
      num_train_steps=num_train_steps,
      num_warmup_steps=num_warmup_steps,
      use_tpu=FLAGS.use_tpu,
      use_one_hot_embeddings=FLAGS.use_tpu)

  estimator = tf.contrib.tpu.TPUEstimator(
      use_tpu=FLAGS.use_tpu,
      model_fn=model_fn,
      config=run_config,
      train_batch_size=FLAGS.train_batch_size,
      eval_batch_size=FLAGS.eval_batch_size,
      predict_batch_size=FLAGS.predict_batch_size)

初始化 model_fn 和估计函数，如果不存在TPU，那么就会初始化GPU或者CPU的估计函数。

2.4.2 训练(train)

  if FLAGS.do_train:
    train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
    file_based_convert_examples_to_features(
        train_examples, label_list, FLAGS.max_seq_length, tokenizer, train_file)
    tf.logging.info("***** Running training *****")
    tf.logging.info("  Num examples = %d", len(train_examples))
    tf.logging.info("  Batch size = %d", FLAGS.train_batch_size)
    tf.logging.info("  Num steps = %d", num_train_steps)
    train_input_fn = file_based_input_fn_builder(
        input_file=train_file,
        seq_length=FLAGS.max_seq_length,
        is_training=True,
        drop_remainder=True)
    estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)

如果需要执行训练，那么就执行此部分代码，这部分代码的作用就是获取和 tfrecord 相关的输入函数然后采用估计函数中的 train 方法来进行训练。

2.4.3 验证(eval)

  if FLAGS.do_eval:
    eval_examples = processor.get_dev_examples(FLAGS.data_dir)
    num_actual_eval_examples = len(eval_examples)
    if FLAGS.use_tpu:
        eval_examples.append(PaddingInputExample())

首先获取验证集，如果使用TPU那么验证集的大小必须是批大小的整数倍，因此需要对其进行填充。

    eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
    file_based_convert_examples_to_features(
        eval_examples, label_list, FLAGS.max_seq_length, tokenizer, eval_file)

    tf.logging.info("***** Running evaluation *****")
    tf.logging.info("  Num examples = %d (%d actual, %d padding)",
                    len(eval_examples), num_actual_eval_examples,
                    len(eval_examples) - num_actual_eval_examples)
    tf.logging.info("  Batch size = %d", FLAGS.eval_batch_size)

把验证集的数据转换成 tfrecord 格式的数据方便训练，打印验证信息。

    eval_steps = None
    if FLAGS.use_tpu:
      assert len(eval_examples) % FLAGS.eval_batch_size == 0
      eval_steps = int(len(eval_examples) // FLAGS.eval_batch_size)

告诉模型使用全部的验证数据进行验证，如果使用 TPU 还需要告诉模型验证步数的大小

    result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)

    output_eval_file = os.path.join(FLAGS.output_dir, "eval_results.txt")
    with tf.gfile.GFile(output_eval_file, "w") as writer:
      tf.logging.info("***** Eval results *****")
      for key in sorted(result.keys()):
        tf.logging.info("  %s = %s", key, str(result[key]))
        writer.write("%s = %s\n" % (key, str(result[key])))

调用模型的方法获得验证结果，写入文件中并打印验证结果

2.4.4 预测(predict)

  if FLAGS.do_predict:
    predict_examples = processor.get_test_examples(FLAGS.data_dir)
    num_actual_predict_examples = len(predict_examples)
    if FLAGS.use_tpu:
        predict_examples.append(PaddingInputExample())

获取预测集，如果使用TPU对数据进行填充。

    predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
    file_based_convert_examples_to_features(predict_examples, label_list,
                                            FLAGS.max_seq_length, tokenizer,
                                            predict_file)

将预测集的数据转换成 tfrecord 格式的数据。

    tf.logging.info("***** Running prediction*****")
    tf.logging.info("  Num examples = %d (%d actual, %d padding)",
                    len(predict_examples), num_actual_predict_examples,
                    len(predict_examples) - num_actual_predict_examples)
    tf.logging.info("  Batch size = %d", FLAGS.predict_batch_size)

    predict_drop_remainder = True if FLAGS.use_tpu else False
    predict_input_fn = file_based_input_fn_builder(
        input_file=predict_file,
        seq_length=FLAGS.max_seq_length,
        is_training=False,
        drop_remainder=predict_drop_remainder)

    result = estimator.predict(input_fn=predict_input_fn)

打印信息，进行预测，获取结果。

    output_predict_file = os.path.join(FLAGS.output_dir, "test_results.tsv")
    with tf.gfile.GFile(output_predict_file, "w") as writer:
      num_written_lines = 0
      tf.logging.info("***** Predict results *****")
      for (i, prediction) in enumerate(result):
        probabilities = prediction["probabilities"]
        if i >= num_actual_predict_examples:
          break
        output_line = "\t".join(
            str(class_probability)
            for class_probability in probabilities) + "\n"
        writer.write(output_line)
        num_written_lines += 1
    assert num_written_lines == num_actual_predict_examples

保存并打印预测结果。

2.5 main函数

if __name__ == "__main__":
  flags.mark_flag_as_required("data_dir")
  flags.mark_flag_as_required("task_name")
  flags.mark_flag_as_required("vocab_file")
  flags.mark_flag_as_required("bert_config_file")
  flags.mark_flag_as_required("output_dir")
  tf.app.run()

设定必须的参数，运行程序。

3 总结

通过对 run_classifier 进行分析，我们可以基本了解如何利用 bert 模型进行分类任务，如何根据自己的任务对 BERT 模型代码进行微调。

下一篇可能是对BERT模型的阅读或者是其他用法的解读 …… 吧。