酸洗嵌套函数的替代方法

Alternative to nested functions for pickling

我有一段代码可以从许多较小的函数生成一个函数,同时让最外层的函数接受一个参数 x

换句话说,我有一个输入 x,我需要对运行时决定的 x 进行各种转换。

这是通过迭代调用这个函数来完成的(它本质上是将一个函数包装在另一个函数中)。

函数如下:

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

    def _function(x):
        return curr_layer(prev_layer(x) if callable(prev_layer) else x)

    return _function

旁注:如您所见,如果 prev_layer 不可调用,它会替换为输入 x 所以我使用虚拟整数来指示输入的位置去了。

问题: 这段代码无法 pickle。 我似乎无法想出一种方法来重写这段代码,使其成为可腌制的。

注意:我需要将此对象持久保存在磁盘上,但也需要将其用于为 IPC 腌制的多处理中(这些函数在那里不使用,因此从技术上讲它们可以感动)

我还有一个处理多个输入的函数的更复杂版本(使用固定聚合函数,在本例中 torch.cat)我知道这两个可以合并为一个通用函数,我会这样做一旦我开始工作。

这是第二个函数的代码:

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

    def _function(x):
        return curr_layer(torch.cat([layer(x) if callable(layer) else x for layer in prev_layers]))

    return _function

我通过将这些函数的 return 值附加到 class 实例来解决这个问题,如 this thread.

中所述