TensorFlow 2.0 中 lambda 函数的组合

Composition of lambda functions in TensorFlow 2.0

为了在 TensorFlow 2.0 中构建一个大模型,我使用了函数式方法并利用了 Python 3.6 的 functools 模块。我通过显示特定自定义层的代码来说明问题。

import tensorflow as tf
import functools
from collections.abc import Iterable


# TODO Check for correctness of the model implementation
class Unit3D(tf.keras.layers.Layer):
    def __init__(self, output_channels,
                 kernel_shape=(1, 1, 1),
                 stride=(1, 1, 1),
                 activation_fn='relu',
                 use_batch_norm=True,
                 use_bias=False,
                 is_training=False,
                 name='unit_3d'):
        super(Unit3D, self).__init__(name=name)
        self._output_channels = output_channels
        self._kernel_shape = kernel_shape
        self._stride = stride
        self._activation = activation_fn
        self._use_batch_norm = use_batch_norm
        self._use_bias = use_bias
        self._is_training = is_training
        self._pipeline = []
        self._pipeline.append(tf.keras.layers.Conv3D(
            filters=self._output_channels,
            kernel_size=self._kernel_shape,
            strides=self._stride,
            padding='same',
            use_bias=self._use_bias,
            data_format='channels_first'
        )
        )
        if self._use_batch_norm:
            bn = tf.keras.layers.BatchNormalization(
                axis=1,
                fused=False,
            )
            bn = functools.partial(bn, training=self._is_training)
            self._pipeline.append(bn)

        if self._activation is not None:
            self._pipeline.append(tf.keras.layers.Activation(
                activation=self._activation
            )
            )

        print(isinstance(self._pipeline, Iterable))
        print(type(self._pipeline))
        self._pipeline = lambda x: functools.reduce(lambda f, g: g(f), self._pipeline, x)

    def call(self, input):
        return self._pipeline(input)

使用以下代码进行测试时,returns 错误

TypeError: reduce() arg 2 must support iteration

错误与__init__方法中self._pipeline中的函数组合有关。

import tensorflow as tf
from nets.i3d import Unit3D

model = Unit3D(output_channels=64, kernel_shape=[7,7,7],
               is_training=True)

input = tf.keras.backend.random_uniform(shape=(1,3,64,224,224),
                                        dtype=tf.float32)
output = model(input)

在 TensorFlow 2.0 中,在急切执行期间,所有列表都包装在称为 <class 'tensorflow.python.training.tracking.data_structures.ListWrapper'>

的数据结构中

原来这个数据结构是可迭代的。我使用 collections.abc 模块中的 Iterable class 对其进行了测试。

我无法理解此代码的问题,并且不确定这是 TensorFlow 2.0 中的内部问题,还是我在这里遗漏的一些基本问题。

如果有帮助,我正在使用从 r2.0 分支的源代码编译的 tensorflow 版本 2.0.0-beta1。对应的git散列为8e423e3d56390671f0d954c90f4fd163ab02a9c1.

我认为您的问题与 Tensorflow 无关。

这一行:

self._pipeline = lambda x: functools.reduce(lambda f, g: g(f), self._pipeline, x)

您正在用一个函数覆盖您之前在构造函数中创建的列表。因此,在 lambda 的实际执行过程中(尤其是您传递给 reduce 的那个),它只是一个 function.

如果将外部 lambda 重构为用于调试目的的常规函数​​,它会变得更加明显:

    def debug_func(x):
        print(type(self._pipeline))  # prints <class 'function'>
        return functools.reduce(lambda f, g: g(f), self._pipeline, x)

    self._pipeline = debug_func  # Clearly, no longer a list

解决方案

我认为你的意思是将 pipeline-as-function 保存在不同于 _pipeline 的字段中,例如:

    # ...
    self._pipeline_func = lambda x: functools.reduce(lambda f, g: g(f), self._pipeline, x)

def call(self, input):
    return self._pipeline_func(input)