在 keras 層包裝 python 可調用

Question

在 keras / tensorflow 中，將層直接描述為函數通常非常簡單 map 它們的輸入到 output，如下所示：

def resnet_block(x, kernel_size):
   ch = x.shape[-1]
   out = Conv2D(ch, kernel_size, strides = (1,1), padding='same', activation='relu')(x)
   out = Conv2D(ch, kernel_size, strides = (1,1), padding='same', activation='relu')(out)
   out = Add()([x,out])
   return out

而子類化Layer得到類似的東西

r = ResNetBlock(kernel_size=(3,3))
y = r(x)

有點麻煩（對於更復雜的例子甚至更麻煩）。

由於 keras 在第一次被調用時似乎非常樂意構建其層的底層權重，我想知道是否有可能只包裝上面的函數並讓 keras 在有輸入時解決問題，即我希望它看起來像這樣：

r = FunctionWrapperLayer(lambda x:resnet_block(x, kernel_size=(3,3)))
y = r(x)

我嘗試實現FunctionWrapperLayer ，如下所示：

class FunctionWrapperLayer(Layer):
    def __init__(self, fn):
        super(FunctionWrapperLayer, self).__init__()
        self.fn = fn
    
    def build(self, input_shape):
        shape = input_shape[1:]
        inputs = Input(shape)
        outputs = self.fn(inputs)
        self.model = Model(inputs=inputs, outputs=outputs)
        self.model.compile()
        
    def call(self, x):
        return self.model(x)

這看起來可能有效，但是每當我使用激活時，我都會遇到一些奇怪的問題，例如

def bad(x):
    out = tf.keras.activations.sigmoid(x)
    out = Conv2D(1, (1,1), strides=(1,1), padding='same')(out)
    return out
x = tf.constant(tf.reshape(tf.range(48,dtype=tf.float32),[1,4,-1,1])
w = FunctionWrapperLayer(bad)
w(x)

我收到以下錯誤

FailedPreconditionError:  Error while reading resource variable _AnonymousVar34 from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/_AnonymousVar34/class tensorflow::Var does not exist.
     [[node conv2d_6/BiasAdd/ReadVariableOp (defined at <ipython-input-33-fc380d9255c5>:12) ]] [Op:__inference_keras_scratch_graph_353]

這對我來說意味着在構建方法中初始化模型存在一些固有的錯誤。 也許有人對那里可能發生的事情或如何獲得我想要的功能有更好的想法。

更新：如 jr15 所述，當涉及的 function 僅使用 keras 層時，上述內容確實有效。 然而，以下也有效，這讓我有點困惑：

i = Input(x.shape[1:])
o = bad(i)
model = Model(inputs=i, outputs=o)
model(x)

順便說一下， model.submodules收益

(<tensorflow.python.keras.engine.input_layer.InputLayer at 0x219d80c77c0>,
 <tensorflow.python.keras.engine.base_layer.TensorFlowOpLayer at 0x219d7afc820>,
 <tensorflow.python.keras.layers.convolutional.Conv2D at 0x219d7deafa0>)

這意味着當這樣做時，激活會自動變成一個“TensorFlowOpLayer”。

另一個更新：查看原始錯誤消息，激活似乎不是唯一的罪魁禍首。 如果我刪除卷積並使用包裝器，一切正常，我在檢查子模塊時再次發現“TensorFlowOpLayer”。

Answer 1

您的解決方案確實有效！ 您遇到的問題是tf.keras.activations.sigmoid不是圖層，而是普通的 Tensorflow function。要使其正常工作，請改用keras.layers.Activation("sigmoid")(x) 。 對於更一般的情況，如果您想使用一些 Tensorflow function 作為層，您可以將它包裹在 Lambda 層中，如下所示：

out = keras.layers.Lambda(lambda x: tf.some_function(x))(out)

有關詳細信息，請參閱文檔： https://keras.io/api/layers/core_layers/lambda/

Answer 2

使用 Tensorflow 2.4 顯然現在就可以使用了。 子模塊現在顯示“TFOpLambda”層。

對於任何感興趣的人，這里有一些稍微改進的包裝器代碼，它也適用於多輸入模型：

class FunctionWrapperLayer(Layer):
    def __init__(self, fn):
        super(FunctionWrapperLayer, self).__init__()
        self.fn = fn
        
    def build(self, input_shapes):
        super(FunctionWrapperLayer, self).build(input_shapes)
        if type(input_shapes) is list:
            inputs = [Input(shape[1:]) for shape in input_shapes]
        else:
            inputs = Input(input_shapes[1:])
        outputs = self.fn(inputs)
        self.fn_model = Model(inputs=inputs, outputs=outputs)
        self.fn_model.compile()
        
    def call(self, x):
        return self.fn_model(x)