无法使$TPU_NAME环境变量正常工作

Question

我是个新手！我正在尝试在Kaggle内核上从头开始训练BERT模型。无法使BERT run_pretraining.py脚本在TPU上工作。不过，它在CPU上运行良好。我猜问题出在$TPU_NAME环境变量上。

!python run_pretraining.py \
--input_file='gs://xxxxxxxxxx/*' \
--output_dir=/kaggle/working/model/ \
--do_train=True \
--do_eval=True \
--bert_config_file=/kaggle/input/bert-bangla-test-config/config.json \
--train_batch_size=32 \
--max_seq_length=128 \
--max_predictions_per_seq=20 \
--num_train_steps=20 \
--num_warmup_steps=2 \
--learning_rate=2e-5 \
--use_tpu=True \
--tpu_name=$TPU_NAME

Answer 1

如果脚本使用的是tf.distribute.cluster_resolver.TPUClusterResolver() (https://www.tensorflow.org/api_docs/python/tf/distribute/cluster_resolver/TPUClusterResolver)，那么只需不使用任何参数实例化TPUClusterResolver，它就会自动拾取TPU_NAME (https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/tpu/client/client.py#L47)。

Answer 2

好的，我找到了一个菜鸟的解决方案：

运行：

import os
os.environ

从返回的字典中，您可以获得地址。直接复制粘贴就行了。它的格式为'TPU_NAME'：'grpc://xxxxxxx‘。