Keras多GPU模型保存报错：TypeError: can't pickle NotImplementedType objects

你遇到的这个问题确实和Keras的模型存储机制以及多GPU模型的特殊结构有关，同时不同目录的运行差异也有明确的诱因，下面我来逐一拆解并给出解决方案：

问题原因分析 #

1. Keras内部序列化机制的限制
   当你使用multi_gpu_model创建分布式训练模型时，Keras会生成一个模型包装器来处理多GPU之间的参数同步和计算分发。这个包装器内部包含了一些和分布式训练相关的特殊对象（比如未实现的占位符接口、跨GPU通信的临时组件），这些对象是无法被Python的pickle模块序列化的——而Keras的model.save()底层正是依赖pickle来序列化模型的部分组件，因此触发了这个错误。

2. 不同目录运行结果有差异的诱因
   相同代码在不同目录表现不同，大概率是以下几种情况：

   • 缓存残留：Keras会生成__pycache__或临时缓存文件，某些目录的旧缓存可能残留了之前训练的异常状态，干扰了当前模型的序列化流程；
   • 目录权限不足：你要保存模型的/disk3/seaice/目录可能存在读写权限限制，导致模型保存时无法正常写入临时文件，进而掩盖了真实错误，只抛出了pickle相关的异常；
   • 路径依赖差异：不同目录下的环境变量、依赖库加载路径可能有细微差别，导致多GPU模型的内部组件初始化状态不同，刚好在某些目录触发了序列化问题。

可行解决方案 #

方案1：保存原始基础模型（最推荐） #

multi_gpu_model只是原始模型的分布式训练包装，训练过程中所有权重都会同步到原始基础模型上。因此你只需要保存原始模型即可，完全不影响后续使用：

# 先定义原始基础模型
base_model = Sequential()
base_model.add(Dense(2, activation='relu'))
base_model.add(Dense(1, activation='sigmoid'))

# 创建多GPU包装模型用于训练
model = multi_gpu_model(base_model, gpus=4)
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=3, batch_size=500)
scores = model.evaluate(x_test, y_test)

# 保存原始基础模型而非多GPU包装模型
base_model.save('/disk3/seaice/seaice_keras_model2.h5')

方案2：检查并修复目标目录权限 #

确保你的用户对/disk3/seaice/目录有读写权限，执行以下命令排查并修复：

# 查看目录权限
ls -ld /disk3/seaice/
# 如果权限不足，修改权限（根据实际需求调整权限值）
sudo chmod -R 755 /disk3/seaice/

方案3：清理目录缓存 #

删除代码目录和目标保存目录下的__pycache__文件夹，避免旧缓存干扰：

rm -rf ./__pycache__
rm -rf /disk3/seaice/__pycache__

方案4：单独保存模型权重 #

如果只需要保存模型权重，可以改用save_weights方法，绕开完整模型的序列化问题：

# 训练完成后保存权重
model.save_weights('/disk3/seaice/seaice_weights.h5')

# 后续加载时，先重构模型结构再加载权重
base_model = Sequential()
base_model.add(Dense(2, activation='relu'))
base_model.add(Dense(1, activation='sigmoid'))
base_model.load_weights('/disk3/seaice/seaice_weights.h5')

内容的提问来源于stack exchange，提问作者송준석