使用Keras multiple_gpu_model训练时触发“Can't pickle module object”错误

这个问题我之前在多GPU训练任务里也碰到过，本质是Keras用multiple_gpu_model时需要序列化一些对象（比如模型结构、自定义组件），但遇到了没法被pickle处理的模块级对象。结合你的场景，给你几个实用的解决思路：

1. 检查并修复自定义组件的可序列化性 #

如果你的模型里用到了自定义层、自定义损失函数或自定义指标，这些组件必须是可pickle的。最常见的修复方式是：

• 用@tf.keras.utils.register_keras_serializable()装饰器标记自定义组件
• 确保自定义层实现了get_config()方法（用来保存层的参数）

举个自定义层的例子：

import tensorflow as tf

@tf.keras.utils.register_keras_serializable()
class MyCustomConvLayer(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, **kwargs):
        super().__init__(**kwargs)
        self.filters = filters
        self.kernel_size = kernel_size

    def build(self, input_shape):
        self.conv = tf.keras.layers.Conv2D(self.filters, self.kernel_size)
        super().build(input_shape)

    def call(self, inputs):
        return self.conv(inputs)

    def get_config(self):
        config = super().get_config()
        config.update({
            'filters': self.filters,
            'kernel_size': self.kernel_size
        })
        return config

如果是自定义损失/指标，同样用装饰器标记即可：

@tf.keras.utils.register_keras_serializable()
def custom_iou(y_true, y_pred):
    # 你的IoU计算逻辑
    ...

2. 改用TensorFlow官方推荐的分布式策略 #

multiple_gpu_model其实是比较旧的API，现在TensorFlow更推荐用tf.distribute.MirroredStrategy，它对多GPU训练的支持更稳定，还能避免很多pickle相关的问题。

给你适配你的场景的示例代码：

import tensorflow as tf

# 初始化多GPU分布式策略
strategy = tf.distribute.MirroredStrategy()
print(f"使用 {strategy.num_replicas_in_sync} 块GPU")

# 所有模型构建、编译操作都要放在strategy.scope()里
with strategy.scope():
    # 这里替换成你原来的模型构建代码
    def build_model():
        inputs = tf.keras.Input(shape=(你的输入形状))
        # ... 你的模型层定义 ...
        outputs = tf.keras.layers.Dense(你的输出维度)(x)
        return tf.keras.Model(inputs, outputs)
    
    model = build_model()
    # 编译模型，用你原来的优化器、损失和指标
    model.compile(
        optimizer='adam',
        loss='你的损失函数',
        metrics=['mean_iou']
    )

# 把训练数据转成tf.data.Dataset（比直接传numpy数组更适合分布式训练）
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(128)
val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(128)

# 开始训练
model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=100
)

3. 排查回调函数中的不可序列化对象 #

如果你在训练时用了自定义回调函数，也要检查回调里有没有引用模块级别的对象（比如某个第三方库的全局实例）。如果有，要么把这些对象改成可序列化的，要么在回调里动态创建它们，而不是直接引用外部模块对象。

4. 手动定位无法序列化的对象 #

如果上面的方法还没解决，可以用pickle手动测试模型相关组件，定位具体是哪个对象出了问题：

import pickle

# 测试模型是否可序列化
try:
    pickle.dumps(model)
    print("模型可正常序列化")
except Exception as e:
    print(f"序列化模型时出错: {e}")

# 测试损失函数/指标
try:
    pickle.dumps(model.loss)
    print("损失函数可正常序列化")
except Exception as e:
    print(f"序列化损失函数时出错: {e}")

根据报错信息，就能精准找到需要修复的对象了。

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