如何构建具有未知输入形状的自定义 keras 层

Question

我正在尝试使用以下代码在 keras 中构建自定义 ConvLSTM 层，但它不起作用：

import tensorflow as tf
from tensorflow import keras
from keras.layers import InputSpec, Layer 

class Padding2D(Layer):
    def __init__(self, padding = (1,1), **kwargs):
        self.padding = tuple(padding)
        self.input_spec = [InputSpec(ndim = 4)]
        super(Padding2D,self).__init__(**kwargs)

    def compute_output_shape(self, s):
        return (s[0], s[1] + 2*self.padding[0], s[2] + 2*self.padding[1], s[3])

    def call(self, x):
        w_pad, h_pad = self.padding 
        return tf.pad(x, [[0,0], [h_pad,h_pad],[w_pad,w_pad],[0,0]])

class ConvLSTM(Layer):
    def __init__(self, out_channels, kernel_size=5, forget_bias=1.0, padding=0):
        super(ConvLSTM, self).__init__()
        self.out_channels = out_channels
        self.forget_bias = forget_bias
        self.states = None

    def call(self, inputs):
        if self.states is None:
            #inputs.shape : [Batch, Height, Width, Channel]
            self.states = (tf.zeros([inputs.shape[0], inputs.shape[1], inputs.shape[2], self.out_channels]),
                  tf.zeros([inputs.shape[0], inputs.shape[1], inputs.shape[2]], self.out_channels))

        c, h = self.states
        if not (len(c.shape) == 4 and len(h.shape) == 4 and len(inputs.shape) == 4):
            raise TypeError("Incorrect shapes")

        inputs_h = tf.concat((inputs, h), axis=3)
        padded_inputs_h = Padding2D(padding = (padding,padding))(inputs_h)
        i_j_f_o = Conv2D( 4 * out_channels, kernel_size, strides=1)(padded_inputs_h)
        i = i_j_f_o[:,:,:,: self.out_channels]
        j = i_j_f_o[:,:,:,self.out_channels : 2*self.out_channels]
        f= i_j_f_o[:,:,:, 2*self.out_channels : 3*self.out_channels]
        o = i_j_f_o[:,:,:, 3*self.out_channels :]
        # i, j, f, o = torch.split(i_j_f_o,  self.out_channels, dim=3)

        new_c = c * sigmoid(f + self.forget_bias) + sigmoid(i) * tanh(j)
        new_h = tanh(new_c) * sigmoid(o)
        self.states = (new_c, new_h)

        return new_h

input0 = tf.keras.Input(shape= (2,2,1))
x = ConvLSTM(out_channels= 1)(input0)
model = tf.keras.Model(input0,x)
print(model(tf.ones((1,2,2,1))))

错误输出

----> x = ConvLSTM(out_channels= 1)(input0)
TypeError: in user code:

    <ipython-input-1-2e11c0026581>:28 call  *
        self.states = (tf.zeros([inputs.shape[0], inputs.shape[1], inputs.shape[2], self.out_channels]),
TypeError: Expected int32, got None of type 'NoneType' instead.

我认为发生错误是因为模型事先不知道在模型构建时（执行前）设置为 None 的 batch_size 维度（inputs.shape[0]）的值，但我需要制作模型图在执行时自行输出批量大小维度（并在构建时忽略它）。 有人可以帮忙吗？

Answer 1

按照上面 Marc 在评论中给出的建议，这段代码解决了这个问题：

import tensorflow as tf
from tensorflow import keras
from keras.layers import InputSpec, Layer, Conv2D 
from tensorflow.keras.activations import sigmoid, tanh

class Padding2D(Layer):
    def __init__(self, padding = (1,1), **kwargs):
        self.padding = tuple(padding)
        self.input_spec = [InputSpec(ndim = 4)]
        super(Padding2D,self).__init__(**kwargs)

    def compute_output_shape(self, s):
        return (s[0], s[1] + 2*self.padding[0], s[2] + 2*self.padding[1], s[3])

    def call(self, x):
        w_pad, h_pad = self.padding 
        return tf.pad(x, [[0,0], [h_pad,h_pad],[w_pad,w_pad],[0,0]])

class ConvLSTM(Layer):
    def __init__(self, out_channels, kernel_size=1, forget_bias=1.0, padding=0):
        super(ConvLSTM, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.forget_bias = forget_bias
        self.padding=padding
        self.states = None

    def call(self, inputs):
        if self.states is None:
            #inputs.shape : [Batch, Height, Width, Channel]
            self.states = (  tf.zeros_like(tf.tile(tf.expand_dims(inputs[:,:,:,0], axis=-1), (1,1,1,self.out_channels))),
               tf.zeros_like(tf.tile(tf.expand_dims(inputs[:,:,:,0], axis=-1), (1,1,1,self.out_channels))))
    
        c, h = self.states
        if not (len(c.shape) == 4 and len(h.shape) == 4 and len(inputs.shape) == 4):
            raise TypeError("Incorrect shapes")

        inputs_h = tf.concat((inputs, h), axis=3)
        padded_inputs_h = Padding2D(padding = (self.padding,self.padding))(inputs_h)
        i_j_f_o = Conv2D( 4 * self.out_channels, self.kernel_size, strides=1)(padded_inputs_h)
        i = i_j_f_o[:,:,:,: self.out_channels]
        j = i_j_f_o[:,:,:,self.out_channels : 2*self.out_channels]
        f= i_j_f_o[:,:,:, 2*self.out_channels : 3*self.out_channels]
        o = i_j_f_o[:,:,:, 3*self.out_channels :]

        new_c = c * sigmoid(f + self.forget_bias) + sigmoid(i) * tanh(j)
        new_h = tanh(new_c) * sigmoid(o)
        self.states = (new_c, new_h)

        return new_h

我还找到了另一种解决问题的方法，即在初始化图层时提供批量大小和输入形状。 代码如下：

import tensorflow as tf
from tensorflow import keras
from keras.layers import InputSpec, Layer, Conv2D 
from tensorflow.keras.activations import sigmoid, tanh

class Padding2D(Layer):
    def __init__(self, padding = (1,1), **kwargs):
        self.padding = tuple(padding)
        self.input_spec = [InputSpec(ndim = 4)]
        super(Padding2D,self).__init__(**kwargs)

    def compute_output_shape(self, s):
        return (s[0], s[1] + 2*self.padding[0], s[2] + 2*self.padding[1], s[3])

    def call(self, x):
        w_pad, h_pad = self.padding 
        return tf.pad(x, [[0,0], [h_pad,h_pad],[w_pad,w_pad],[0,0]])

class ConvLSTM(Layer):
    def __init__(self,batch_size, input_shape, out_channels, kernel_size=1, forget_bias=1.0, padding=0):
        super(ConvLSTM, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.forget_bias = forget_bias
        self.shape = input_shape
        self.padding=padding
        self.states = None
        self.batch_size = batch_size

    def build(self, input_shape):
        if self.states is None:
            #input_shape : [Height, Width, Channel]  
            self.states = (tf.zeros([self.batch_size]+ self.shape[:-1] + [self.out_channels]),
               tf.zeros([self.batch_size]+ self.shape[:-1] + [self.out_channels]))
        super(ConvLSTM,self).build(input_shape)

    def call(self, inputs):
        c, h = self.states
        if not (len(c.shape) == 4 and len(h.shape) == 4 and len(inputs.shape) == 4):
            raise TypeError("Incorrect shapes")

        inputs_h = tf.concat((inputs, h), axis=3)
        padded_inputs_h = Padding2D(padding = (self.padding,self.padding))(inputs_h)
        i_j_f_o = Conv2D( 4 * self.out_channels, self.kernel_size, strides=1)(padded_inputs_h)
        i,j,f,o = tf.split(i_j_f_o, num_or_size_splits=4, axis=3)

        new_c = c * sigmoid(f + self.forget_bias) + sigmoid(i) * tanh(j)
        new_h = tanh(new_c) * sigmoid(o)
        self.states = (new_c, new_h)

        return new_h

然而，即使这些实现解决了本文中提出的问题，在与我如何更新 lstm 单元状态相关的两个实现中仍然存在问题（行 self.states = (new_c, new_h) """ConvLSTM 类中的最后一行” """) 但由于问题不同，我在不同的帖子中打开了这个问题