使用Kaggle TPU训练TensorFlow模型保存时遇UnimplementedError求助

报错原因 #

Kaggle TPU处于分布式运行环境，直接在TPU策略上下文内保存模型时，文件系统操作会绑定到TPU的分布式存储体系，导致本地路径的文件系统操作不被支持，触发UnimplementedError。

解决方案 #

以下几种方法均可解决该问题：

1. 切换到CPU上下文保存模型

将模型保存操作放在CPU设备上下文中执行，绕开TPU的文件系统限制：

import tensorflow as tf

# 切换至CPU设备执行保存操作
with tf.device('/CPU:0'):
    model_sub.save('saved_model/xception')

2. 先保存权重再在CPU上重建模型保存

如果直接切换设备仍有问题，可先保存模型权重，再在CPU环境下重建模型并加载权重后保存：

# 先保存模型权重
model_sub.save_weights('xception_weights.h5')

# 在CPU上下文重建模型、加载权重并保存
with tf.device('/CPU:0'):
    # 重新构建模型结构
    model_cpu = keras.models.Sequential(
        [keras.layers.Input([None, None, 3])] + 
        model.layers[:-1] + 
        [tf.keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=-1), name='embedding_norm')]
    )
    # 加载预存权重
    model_cpu.load_weights('xception_weights.h5')
    # 保存完整模型
    model_cpu.save('saved_model/xception')

3. 指定Kaggle可写目录保存

Kaggle环境中，/kaggle/working/是默认的可写目录，直接指定该路径也可避免文件系统权限问题：

model_sub.save('/kaggle/working/saved_model/xception')

内容的提问来源于stack exchange，提问作者Youseef Moemen