时间戳

首页 > 解决方案 > 无法让 $TPU_NAME 环境变量正常工作

tensorflow - 无法让 $TPU_NAME 环境变量正常工作

问题描述

我是新手！我正在尝试在 Kaggle 内核上从头开始训练 BERT 模型。无法让 BERTrun_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

标签： tensorflowkagglebert-language-modelgoogle-cloud-tpu

解决方案

如果脚本使用的是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）。

推荐阅读