如何在Kaggle Cloud TPU上正确使用TensorFlow Hub模型？

阿华AIGC实验室

2026-4-30

解决Kaggle TPU环境下TensorFlow Hub模型加载的文件系统错误

先还原我的问题场景：
我在Kaggle平台上用TensorFlow Hub的MobileNetV2特征向量构建分类模型，代码如下：

m = tf.keras.Sequential([
    hub.KerasLayer("https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4", output_shape=[1280], trainable=False),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])
m.build([None, 224, 224, 3]) # Batch input shape.

这段代码在GPU环境下运行完全正常，但切换到TPU环境并使用TFRecords数据集时，突然抛出了这个错误：

InvalidArgumentError: Unsuccessful TensorSliceReader constructor: Failed to get matching files on /tmp/tfhub_modules/87fb99f72aec02d017e12c0a3d86c5c182ec22ca/variables/variables: Unimplemented: File system scheme '[local]' not implemented (file: '/tmp/tfhub_modules/87fb99f72aec02d017e12c0a3d86c5c182ec22ca/variables/variables')

排查后确认环境配置和TFRecords数据集都没问题——把预训练模型换成Keras官方实现的MobileNetV2后，代码就能正常跑。我试着做了缓存操作但没成功，想请教下遵循TensorFlow Hub缓存指南时，针对Kaggle TPU这种场景有哪些必须注意的点？

核心原因先搞懂

这个错误的本质是TPU节点无法访问本地文件系统：TF Hub默认会把下载的模型缓存到本地的/tmp/tfhub_modules目录，但Kaggle的TPU是独立于你的运行节点的集群，TPU worker根本读不到本地磁盘上的缓存文件；而Keras官方的预训练模型是直接在TPU上初始化权重的，不存在跨节点文件访问的问题，所以能正常运行。

遵循TF Hub缓存指南的关键注意事项（适配Kaggle TPU）

根据我的踩坑经验，结合缓存指南的要求，你需要重点关注这几点：

强制指定TPU可访问的缓存目录
Kaggle的TPU环境允许访问/kaggle/working或者配置后的共享存储路径，你必须通过环境变量TFHUB_CACHE_DIR把缓存目录设置到这些共享路径下。比如在代码最开头就加：
```
import os
os.environ["TFHUB_CACHE_DIR"] = "/kaggle/working/tfhub_cache"
```
这样TF Hub会把模型缓存到/kaggle/working下的子目录，TPU节点可以正常读取这里的文件。
提前在CPU/GPU上下文预加载缓存
一定要在构建TPU策略之前，先在CPU或GPU环境下加载一次TF Hub模型，确保缓存已经生成到指定的共享目录里。比如：
```
import tensorflow as tf
import tensorflow_hub as hub

# 先在CPU环境加载模型，完成缓存写入
with tf.device('/cpu:0'):
    temp_model = tf.keras.Sequential([
        hub.KerasLayer("https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4", trainable=False)
    ])
    temp_model.build([None,224,224,3])
```
这样后续在TPU策略下构建模型时，会直接读取已缓存的文件，不会在TPU节点上触发下载或缓存操作（这正是报错的根源）。
绝对避开本地临时目录
别再用默认的/tmp或者其他本地临时路径了，TPU集群的worker节点完全访问不了这些地方，必须用Kaggle提供的共享存储（/kaggle/input或/kaggle/working）。
提前创建缓存目录并确保权限
虽然Kaggle环境下默认权限没问题，但为了保险起见，可以提前创建缓存目录：
```
os.makedirs(os.environ["TFHUB_CACHE_DIR"], exist_ok=True)
```
避免因为目录不存在导致缓存失败。

用hub.load手动控制模型加载流程
你可以先通过hub.load()加载模型，手动保存到共享路径，再在KerasLayer中使用本地路径，这样流程更可控：

model_handle = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4"
loaded_model = hub.load(model_handle)
# 手动保存到共享目录
tf.saved_model.save(loaded_model, "/kaggle/working/mobilenet_v2_feature_vector")
# 然后用本地路径构建模型
m = tf.keras.Sequential([
    hub.KerasLayer("/kaggle/working/mobilenet_v2_feature_vector", output_shape=[1280], trainable=False),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])