GCP TPU VM运行自定义神经网络时初始化及连接故障求助

问题修复步骤 #

1. 移除干扰性环境变量
   删除代码中os.environ['TPU_LOAD_LIBRARY'] = "0"这一行。TPU VM环境下TensorFlow会自动加载适配的TPU库，手动设置该变量会破坏默认加载逻辑，引发核心转储。

2. 使用标准TPU初始化流程
   替换原有的TPU初始化代码，使用官方适配TPU VM的逻辑：

   import tensorflow as tf
   
   resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
   tf.config.experimental_connect_to_cluster(resolver)
   tf.tpu.experimental.initialize_tpu_system(resolver)
   strategy = tf.distribute.TPUStrategy(resolver)
   

   这段代码会自动识别本地TPU节点，无需手动指定TPU_NAME为local。

3. 确保环境版本匹配
   你创建的TPU VM镜像版本为tpu-vm-tf-2.11.0，因此代码中不要手动升级或降级TensorFlow版本，保持TF 2.11.x即可。可运行以下命令验证TPU硬件连接状态：

   curl http://localhost:8475/requestversion
   

   返回正常版本信息则说明硬件连接无问题。

4. 在TPU策略作用域内构建模型
   所有模型定义、编译、训练的代码必须放在strategy.scope()下，示例：

   with strategy.scope():
       # 自定义神经网络定义
       model = tf.keras.Sequential([
           tf.keras.layers.Dense(64, activation='relu'),
           tf.keras.layers.Dense(10, activation='softmax')
       ])
       model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
   

5. 重启实例排除异常状态
   若以上操作无效，执行命令重启TPU VM：

   gcloud compute tpus tpu-vm restart test-tpu-vm --zone=us-central1-b
   

   重启后重新克隆代码、安装依赖再运行脚本。

内容的提问来源于stack exchange，提问作者Brad Messer