TensorFlow搭配multiprocessing报cannot pickle 'weakref' object错误怎么办

问题根因 #

你遇到的报错不是因为子进程无法使用TensorFlow模型，而是因为Python multiprocessing 模块在Windows平台默认使用spawn模式创建子进程，需要把传递给子进程的所有对象通过pickle序列化，而TensorFlow/Keras的模型实例内部包含弱引用、张量句柄等无法被pickle序列化的对象。你在主进程先完成了模型初始化，再把带模型的Partition实例传给子进程执行test方法，序列化阶段就会直接失败。

解决方法 #

把模型初始化的逻辑放到子进程内部执行，不要在主进程提前创建TensorFlow模型再跨进程传递。修改后的示例代码如下：

import tensorflow as tf
import multiprocessing as mp


class SeqModel(tf.keras.Sequential):
    def __init__(self, input_size, hidden_sizes, output_size):
        super().__init__()
        self.add(tf.keras.layers.Dense(hidden_sizes[0], activation="relu", input_shape=(input_size,)))
        for hidden_size in hidden_sizes[1:]: 
            self.add(tf.keras.layers.Dense(hidden_size, activation="relu"))
        if output_size is not None: 
            self.add(tf.keras.layers.Dense(output_size))


class Partition:
    def __init__(self, partition_id):
        self.partition_id = partition_id
        self.model = None

    def initialization(self):
        self.model = SeqModel(10,[10,10],10)

    def test(self):
        # 模型初始化放到子进程内部执行
        self.initialization()
        print(f'partition {self.partition_id} testing...')


def func():
    # 主进程只传可序列化的基础参数，不传带模型的实例
    partition_ids = list(range(4))
    p = mp.Pool(4)
    for pid in partition_ids:
        # 直接传partition id，在子进程里创建Partition实例、初始化模型
        p.apply(Partition(pid).test)
    p.close()
    p.join()


if __name__ == '__main__':
    func()

如果需要给每个进程的模型传递不同的超参数、训练数据，只要把这些可序列化的基础参数（整数、字符串、普通列表、numpy数组等）从主进程传给子进程，在子进程内部完成模型创建、数据加载、训练全流程即可，完全可以实现多进程独立训练多个模型的需求。

补充说明 #

• 如果后续要回传训练完成的模型参数，不要直接传模型实例，可以在子进程把模型权重转成numpy数组序列化回传到主进程，主进程再加载到对应结构的模型里即可。
• 要是有大量数据需要共享，建议提前把数据存到磁盘，子进程各自读取，避免跨进程传大量数据带来的额外开销。

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