TensorFlow Models：在 Cloud TPU 上微调 BERT 完成句/句对分类任务的完整实战（TF 2.1）

本教程演示如何在 Google Cloud TPU 上训练 BERT（Bidirectional Encoder Representations from Transformers）模型，完成句子和句-句对分类任务（以 GLUE 基准中的 MNLI 为例）。读完本文，你将掌握：如何用 ctpu 工具一键拉起 Compute Engine VM + Cloud TPU 实例、如何用 gsutil 管理训练数据与检查点存储、如何调用仓库中的 run_classifier.py 在 8 核 v3 TPU 上微调 BERT-Large，以及训练完成后如何核训练结果、清理资源以停止计费。

需要说明两点前提：该教程对应的脚本入口是 official/legacy/bert/run_classifier.py，属于仓库中 official/legacy/bert 目录下的 TF 2.x BERT 实现（该目录 README 已声明正在逐步弃用，新任务建议参考 official/nlp 下的新版实现，参见 official/legacy/bert/README.md）；教程中 TPU 镜像锁定为 --tf-version=2.1，与当前仓库代码的 TF 2.x 接口一致，可正常运行。

一、整体流程与涉及的仓库组件

整个 Cloud TPU 微调流程由四部分组成，与仓库源码的对应关系如下：

1. 基础设施：GCP 侧创建 Cloud Storage 桶、Compute Engine VM 和 v3-8 Cloud TPU（全部通过 gcloud/gsutil/ctpu 命令行完成，不涉及仓库代码）；
2. 数据准备：GLUE 数据集（tf_record 格式 + meta_data 文件），教程直接使用官方预先处理好的 gs://cloud-tpu-checkpoints/bert/classification 桶；
3. 训练：run_classifier.py --mode=train_and_eval，内部通过 Keras compile/fit API + DistributionStrategy 跑在 TPU 上；
4. 结果验证与资源清理：检查 Training Summary，然后 ctpu delete + gsutil rm 停止计费。

训练命令中的三个关键 flag（--distribution_strategy、--tpu、--num_gpus 等）由 common_flags.py 的 define_common_bert_flags() 统一定义，其底层调度逻辑在 distribute_utils.py 中实现。

二、搭建 Cloud Storage 与 Compute Engine VM

2.1 创建 GCS 存储桶

1. 打开 Cloud Shell 窗口。
2. 创建项目 ID 变量：

export PROJECT_ID=your-project_id

3. 配置 gcloud 命令行工具使用目标项目：

gcloud config set project ${PROJECT_ID}

4. 创建 Cloud Storage 桶：

gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 -b on gs://your-bucket-name

该桶用于存放训练数据与训练结果（检查点、TensorBoard summaries 都会写入 --model_dir，因此必须使用远程存储而非 VM 本地盘）。参数含义：-c standard 为标准存储类，-l europe-west4 指定地域，-b on 开启版本号（versioning）。

2.2 用 ctpu up 拉起 VM + Cloud TPU

ctpu up --tpu-size=v3-8 \
 --machine-type=n1-standard-8 \
 --zone=europe-west4-a \
 --tf-version=2.1 [optional flags: --project, --name]

• --tpu-size=v3-8：8 核 Cloud TPU v3 片元，是本教程的训练设备；
• --machine-type=n1-standard-8：配套的 8 核 CPU VM（TPU 必须与一个 host VM 配对）；
• --zone：TPU 分区，需与训练数据所在 GCS 地域（europe-west4）匹配以降低延迟；
• --tf-version=2.1：VM 上预装的 TensorFlow 版本。

执行后会打印配置摘要，输入 y 确认、n 取消。命令结束后，shell 提示符会从 username@project 变为 username@tpuname，表示你已进入 VM：

gcloud compute ssh vm-name --zone=europe-west4-a
(vm)$ export TPU_NAME=vm-name

后续所有以 (vm)$ 开头的命令都在 VM 会话中执行。注意 TPU_NAME 即 VM 名称，之后会传给训练脚本的 --tpu flag——在 ctpu 工作流中，TPU 地址就是 VM 主机名，TensorFlow 会自动通过 TPU 服务发现连接，无需手动填 IP。

三、准备数据集

1. 在 VM 上安装仓库依赖：

(vm)$ cd /usr/share/models
(vm)$ sudo pip3 install -r official/requirements.txt

注意 ctpu 镜像中仓库代码位于 /usr/share/models（而非 GitHub 克隆路径），依赖文件是仓库根目录下的 official/requirements.txt。

2. （可选）自行下载 GLUE 数据。本教程直接使用 Google 预先处理好的 GLUE 数据，位于 gs://cloud-tpu-checkpoints/bert/classification，其中包含各任务的 *_train.tf_record、*_eval.tf_record 和 *_meta_data 三类文件，因此可跳过数据转换步骤。

如果你想用自己的数据，官方流程是用 create_finetuning_data.py 把原始文本转成 tf_record（详见 official/legacy/bert/README.md 的 "Fine-tuning" 一节）。run_classifier.py 对输入的要求见 run_classifier.py：--input_meta_data_path 指向一个 JSON 文件，其中至少包含 max_seq_length、train_data_size、eval_data_size、num_labels，可选 label_type（int/float，决定是分类还是回归任务）与 has_sample_weights。

训练/评估数据的 tf.data 管道由 input_pipeline.py 的 create_classifier_dataset() 构建，其解码的字段为 input_ids、input_mask、segment_ids、label_ids（句对任务中 sentence A 与 sentence B 的 token 拼接在一起、用 segment_ids 区分）：

name_to_features = {
    'input_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
    'input_mask': tf.io.FixedLenFeature([seq_length], tf.int64),
    'segment_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
    'label_ids': tf.io.FixedLenFeature([], label_type),
}

从源码可以看到两个 TPU 相关的关键细节：

• int64 转 int32：decode_record() 会把所有 int64 特征 tf.cast 成 int32，因为 tf.Example 只支持 int64 而 TPU 只支持 int32（input_pipeline.py）；
• 按 host 分片：create_classifier_dataset() 中 dataset.shard(input_pipeline_context.num_input_pipelines, ...) 按 host 数（而非核数）对数据分片，保证每个 TPU host 读取不同的数据切片（input_pipeline.py）。训练时还会 shuffle(100) 并 repeat()，batch(drop_remainder=is_training) 则保证训练 batch 可被 8 核整除。

四、定义训练参数

在 VM 中定义训练/评估所需的环境变量：

(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/tpu/models"
(vm)$ export STORAGE_BUCKET=gs://your-bucket-name
(vm)$ export BERT_BASE_DIR=gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16
(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/bert-output
(vm)$ export GLUE_DIR=gs://cloud-tpu-checkpoints/bert/classification
(vm)$ export TASK=mnli

各变量含义：

变量	说明
PYTHONPATH	保证 import official... 能解析到 ctpu 镜像中的 /usr/share/tpu/models 仓库代码
STORAGE_BUCKET	你的 GCS 桶，训练输出根目录
BERT_BASE_DIR	BERT-Large（uncased_L-24_H-1024_A-16，24 层/1024 隐层/16 头/340M 参数）预训练检查点目录，含 bert_config.json、bert_model.ckpt、vocab.txt
MODEL_DIR	检查点与 summaries 输出目录
GLUE_DIR	预先处理好的 GLUE 数据目录
TASK	任务名，本教程用 mnli；句/句对分类任务（MRPC、QNLI、QQP 等）均可替换此值

预训练检查点的完整清单（BERT-Base/Large、Cased/Uncased/多语言/Whole Word Masking 共 7 个版本）见 official/legacy/bert/README.md 的 "Pre-trained Models" 一节，其中也建议把检查点托管在 GCS 上供 Cloud GPU/TPU 使用。

五、启动训练

(vm)$ python3 official/nlp/bert/run_classifier.py \
  --mode='train_and_eval' \
  --input_meta_data_path=${GLUE_DIR}/${TASK}_meta_data \
  --train_data_path=${GLUE_DIR}/${TASK}_train.tf_record \
  --eval_data_path=${GLUE_DIR}/${TASK}_eval.tf_record \
  --bert_config_file=$BERT_BASE_DIR/bert_config.json \
  --init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \
  --train_batch_size=32 \
  --eval_batch_size=32 \
  --learning_rate=2e-5 \
  --num_train_epochs=3 \
  --model_dir=${MODEL_DIR} \
  --distribution_strategy=tpu \
  --tpu=${TPU_NAME}

说明：教程原文写的是 official/nlp/bert/run_classifier.py，对应的是 TF 2.1 时期的目录结构；在当前仓库中该脚本位于 official/legacy/bert/run_classifier.py，在 ctpu 的 TF 2.1 镜像中路径保持为 official/nlp/bert/run_classifier.py。

5.1 参数详解（结合源码默认值）

• --mode：枚举值 train_and_eval（默认）、export_only、predict（run_classifier.py）。train_and_eval 边训练边评估；export_only 取 model_dir 中最新检查点导出 SavedModel；predict 恢复检查点后对测试集输出每样本类别概率，结果写入 model_dir/test_results.tsv。
• --train_batch_size=32：全局 batch size（默认 32）。注意在 TPU 下它是跨 8 核的全局批次，create_classifier_dataset 内部通过 ctx.get_per_replica_batch_size(global_batch_size) 均分到每个核（run_classifier.py）。
• --learning_rate=2e-5：初始学习率，flag 默认值是 5e-5（common_flags.py），本教程对 24 层大模型取更小的 2e-5。
• --num_train_epochs=3：默认 3 个 epoch。
• --distribution_strategy=tpu + --tpu=${TPU_NAME}：切换到 TPU 策略。这两个 flag 的解析入口是 distribute_utils.py 的 get_distribution_strategy()：当策略为 tpu 时，构造 TPUCusterResolver、调用 tf.tpu.experimental.initialize_tpu_system() 初始化拓扑，再返回 tf.distribute.TPUStrategy；文档注释明确说明 distribution_strategy=tpu 时 tpu_address 不能为 None，否则抛出 ValueError。
• --init_checkpoint：预训练 BERT 权重。run_keras_compile_fit() 中通过 tf.train.Checkpoint(model=sub_model).read(init_checkpoint) 恢复编码器权重，再用 assert_existing_objects_matched() 校验变量匹配（run_classifier.py）。也可以改用 --hub_module_url 指向 TF-Hub 上的 BERT 模块（二选一，见 official/legacy/bert/README.md）。

5.2 训练循环内部发生了什么

custom_main() 在 mode=train_and_eval 时最终调用 run_bert()，其中几个值得注意的计算（run_classifier.py）：

epochs = FLAGS.num_train_epochs * FLAGS.num_eval_per_epoch
steps_per_epoch = int(train_data_size / FLAGS.train_batch_size)
warmup_steps = int(epochs * train_data_size * 0.1 / FLAGS.train_batch_size)

即 warmup 步数 = 总训练步数的 10%，之后线性衰减到 --end_lr（默认 0）。优化器由 optimization.py 的 create_optimizer() 创建：默认 AdamWeightDecay（--optimizer_type=adamw，也支持 lamb），weight_decay_rate=0.01，学习率调度是 PolynomialDecay 外包一层 WarmUp。

模型构建、编译与拟合流程（run_classifier.py）：

1. 在 strategy.scope() 内创建训练/评估 Dataset 和分类器模型 bert_models.classifier_model()；
2. 分类任务用自定义交叉熵 loss（log_softmax + one-hot，等价于 SparseCategoricalCrossEntropy），num_labels==1 时退化为回归（MeanSquaredError）；分类任务的评估指标是 SparseCategoricalAccuracy；
3. bert_model.compile(..., steps_per_execution=steps_per_loop) + fit()；steps_per_execution 控制多少个训练步打包进同一个 tf.function 执行循环——official/legacy/bert/README.md 建议 TPU 上设 --steps_per_loop=1000，可显著提升 TPU 利用率（代价是循环内不触发 callbacks）；
4. 通过 tf.train.CheckpointManager（max_to_keep=None，即保留所有检查点）+ TensorBoard callback 把检查点与 summaries 写回 MODEL_DIR（GCS）。

训练结束后的返回值正是教程中展示的结果来源：

stats = {'total_training_steps': steps_per_epoch * epochs}
stats['train_loss'] = history.history['loss'][-1]
stats['eval_metrics'] = history.history['val_accuracy'][-1]

六、验证训练结果

在 v3-8 TPU 上，本教程（MNLI，BERT-Large，3 epoch）训练大约需要 1 小时。脚本完成时应看到如下形式的汇总（run_classifier.py 的 stats 字典直接输出）：

Training Summary:
{'train_loss': 0.28142181038856506,
'last_train_metrics': 0.9467429518699646,
'eval_metrics': 0.8599063158035278,
'total_training_steps': 36813}

• train_loss：最后一个 batch 的训练交叉熵；
• eval_metrics：验证集（MNLI-m）accuracy，约 0.86；
• total_training_steps：steps_per_epoch × epochs，即 36813 步。

检查点保存在 ${STORAGE_BUCKET}/bert-output 中，summaries 在其 summaries/ 子目录，可用 TensorBoard 查看。如需对测试集推理，把 --mode 换成 predict 并提供 --eval_data_path（测试集 tf_record），每行类别概率会写入 test_results.tsv（run_classifier.py）。

七、清理资源（停止计费）

为避免 GCP 账单产生不必要费用，按顺序执行：

1. 断开 VM 会话：

(vm)$ exit

2. 在 Cloud Shell 中用创建时相同的 --zone 删除 VM 与 Cloud TPU：

$ ctpu delete --zone=your-zone

3. 确认没有已分配实例（删除可能需要几分钟）：

$ ctpu status --zone=your-zone

4. 删除本教程创建的存储桶（替换为实际桶名）：

$ gsutil rm -r gs://your-bucket

八、关键要点回顾

• 设备侧：ctpu up 一次完成 VM + v3-8 TPU 的拉起，--tf-version=2.1 决定运行时版本；--tpu 传 VM 名即可，TPUStrategy 会自动发现集群（distribute_utils.py）；
• 数据侧：GLUE 数据使用官方 gs://cloud-tpu-checkpoints/bert/classification，输入为 tf_record + JSON meta_data，解码字段与 int32 转换逻辑见 input_pipeline.py；
• 训练侧：全局 batch 32、学习率 2e-5、3 epoch、AdamW + 10% warmup 的线性衰减，检查点持续写回 GCS（run_classifier.py、optimization.py）；
• 结果侧：MNLI 验证精度约 0.86，约 1 小时完成；
• 成本侧：ctpu delete + gsutil rm -r 收尾。

如果要在 GPU 上复现同一流程，只需把 --distribution_strategy=tpu --tpu=${TPU_NAME} 替换为 --distribution_strategy=mirrored，其余参数不变（参考 official/legacy/bert/README.md 中 MRPC 的 GPU 示例）。