Google Colab TPU突然无法识别问题求助（附环境与代码）

问题概述 #

自2025年3月10日起，Colab Pro订阅用户的TPU无法被识别，已尝试切换运行时类型、重启会话、重新登录账号，问题仍未解决。当前环境：

• TensorFlow 2.18.1
• Keras-Tuner 1.4.7
• 可用计算单元55.47，使用率约3.66/小时，存在4个活跃会话

运行初始化代码后输出：

Connecting to TPU...
TPU not found. Falling back to CPU/GPU strategy.
Using distribution strategy: _DefaultDistributionStrategy
Number of replicas: 1
Mixed Precision enabled: <DTypePolicy "mixed_bfloat16">

排查与解决方案 #

1. 确认运行时配置生效 #

• 打开Colab笔记本，依次点击 Runtime > Change runtime type
• 在Hardware accelerator下拉菜单中选择TPU，点击Save
• 再点击 Runtime > Restart runtime，重启后重新运行初始化代码

2. 清理闲置活跃会话 #

Colab的并发会话会占用资源，关闭闲置会话释放TPU配额：

• 进入Colab会话管理界面，关闭所有不需要的活跃会话
• 重启当前笔记本的运行时，再次尝试连接TPU

3. 调整TensorFlow版本适配TPU #

部分TF版本与Colab TPU存在兼容性问题，建议尝试切换到稳定兼容版本：

!pip install tensorflow==2.17.0 --upgrade

运行上述命令后，点击 Runtime > Restart runtime，再重新执行TPU初始化代码

4. 修改TPU初始化代码（显式指定TPU地址） #

默认的TPUClusterResolver可能无法自动识别地址，尝试显式指定：

import os
import tensorflow as tf
from tensorflow.keras import mixed_precision

# TPU Strategy Setup
try:
    print("Connecting to TPU...")
    # 从环境变量获取TPU地址并显式指定
    tpu_addr = 'grpc://' + os.environ['COLAB_TPU_ADDR']
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_addr)
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    print("TPU initialized successfully!")
    strategy = tf.distribute.TPUStrategy(tpu)
except ValueError as e:
    print(f"TPU connection error details: {e}")
    print("Falling back to CPU/GPU strategy.")
    strategy = tf.distribute.get_strategy()

# Confirm strategy definition
if 'strategy' not in locals():
    raise RuntimeError("TPU strategy setup failed. Please check TPU initialization.")

# Log TPU Configuration
print(f"Using distribution strategy: {strategy.__class__.__name__}")
print(f"Number of replicas: {strategy.num_replicas_in_sync}")

tf.debugging.set_log_device_placement(True)

# Enable Mixed Precision for TPU optimization
policy = mixed_precision.Policy('mixed_bfloat16')
mixed_precision.set_global_policy(policy)
print("Mixed Precision enabled: ", policy)

5. 检查Colab服务状态 #

若上述方法均无效，可能是Google端TPU服务临时故障，可查看Colab官方状态公告确认是否存在全局服务异常。

内容的提问来源于stack exchange，提问作者Albert Blarer