寻求替代嵌套函数的实现方案，使构建的动态函数支持pickle序列化

核心问题 #

你当前代码生成的函数无法被pickle序列化的根本原因是使用了动态定义的嵌套闭包函数。pickle序列化函数时要求函数在模块顶层可被导入，动态生成的内部闭包函数不满足该要求，因此序列化失败。

改写方案 #

采用顶层定义的可调用类（实现__call__方法）替代嵌套闭包，类本身是模块顶层可导入的，只要成员变量可序列化，类实例就可正常被pickle。

单输入版本实现 #

import typing
from typing import Union

class LayerWrapper:
    def __init__(self, curr_layer: typing.Callable, prev_layer: Union[typing.Callable, int]):
        self.curr_layer = curr_layer
        self.prev_layer = prev_layer
    
    def __call__(self, x):
        return self.curr_layer(self.prev_layer(x) if callable(self.prev_layer) else x)

def build_layer(curr_layer: typing.Callable, prev_layer: Union[typing.Callable, int]) -> typing.Callable:
    return LayerWrapper(curr_layer, prev_layer)

多输入版本实现 #

import typing
import torch

class MultiInputLayerWrapper:
    def __init__(self, curr_layer: typing.Callable, prev_layers: list):
        self.curr_layer = curr_layer
        self.prev_layers = prev_layers
    
    def __call__(self, x):
        return self.curr_layer(torch.cat([layer(x) if callable(layer) else x for layer in self.prev_layers]))

def build_layer_multi_input(curr_layer: typing.Callable, prev_layers: list) -> typing.Callable:
    return MultiInputLayerWrapper(curr_layer, prev_layers)

效果验证 #

可以通过以下代码测试序列化能力：

import pickle

# 测试用自定义层
def add_one(x):
    return x + 1
def multiply_two(x):
    return x * 2

# 构建多层函数
layer1 = build_layer(add_one, 0) # 0为虚拟值标识输入位
layer2 = build_layer(multiply_two, layer1)

# 原始调用测试，预期输出 (3+1)*2=8
print(layer2(3))

# 序列化&反序列化测试
pickled_data = pickle.dumps(layer2)
restored_layer = pickle.loads(pickled_data)
# 反序列化后调用测试，同样输出8
print(restored_layer(3))

注意事项 #

• 只要你传入的curr_layer和prev_layer本身支持pickle（比如常规PyTorch层、顶层自定义函数），整个生成的可调用对象就可正常序列化、反序列化
• 后续要合并为通用实现只需要调整类的初始化和__call__逻辑即可，序列化能力不受影响
• 完全适配你的多进程IPC场景，哪怕反序列化后不调用函数，也可正常完成数据传输

内容的提问来源于stack exchange，提问作者Igor Vereš