您可以将keras.layers输出与TensorFlow操作混合吗？

Question

我知道Keras层的输出（例如keras.layers.dense（））产生所谓的“ keras张量”。此外，还有由TensorFlow操作产生的“ TensorFlow Tensors”（例如tf.Reduce_sum（））。

某些功能只能通过其中一种方法传递，因此很明显，有时我必须将它们混合才能进行计算。我认为，将Keras层代码与TF OPS混合看起来远非美丽，但这就是另一个话题。

我的问题是 - 混合凯拉斯层和张量吗？并且是否存在与将Keras和Tensorflow张量混合有关的问题？

---

让我们考虑一个自定义类的示例，不要深入研究计算，这是没有意义的。关键是要展示我在说什么。

class Calculation(layers.Layer):
    def __init__(self, units):
        super().__init__()
        self.units = units
        self.conv3d = layers.Conv3D(units, kernel_size=3, strides=2)
        self.norm = layers.LayerNormalization(units)
    
    def call(self, x):
        x = activations.gelu(x)
        x = self.norm(x)  # keras layer
        x = tf.transpose(x, (0, 4, 1, 2, 3))
        x = self.conv3d(x)  # keras layer
        x = tf.transpose(x, (0, 2, 3, 4, 1))
        x = tf.reshape(x, (-1, 1, 2 * self.units))
    
        return x

---

另一个示例，这次不是在自定义层中：

encoder_input = keras.Input(shape=(28, 28, 1), name="img")
x = layers.Conv2D(16, 3, activation="relu")(encoder_input)
x = layers.Conv2D(32, 3, activation="relu")(x)
x = layers.MaxPooling2D(3)(x)
x = layers.Conv2D(32, 3, activation="relu")(x)
x = tf.einsum('...ijk->...jik', x)  # tensorflow op, just swap height and width for no reason
x = layers.Conv2D(16, 3, activation="relu")(x)

---

我知道有keras.layers.lambda layer之类的东西，但是有时我需要使用TensorFlow的 90％ OPS和仅 10％按KERAS层进行的计算（因为没有TF OP替代方案，我不想每次都能实现自己的）。在这种情况下，使用lambda层毫无意义。

在TensorFlow 2.x中编写复杂模型（实现不仅仅堆叠Keras现有层）是否有一个好的方法？

Answer 1

IMO，只要您保留批处理尺寸和张量形状，就可以将Tensorflow操作与keras层混合在一起并不重要。例如，您可能需要将tf在lambda层内包装操作以设置一些元数据，但这取决于您的口味：

import tensorflow as tf

encoder_input = tf.keras.Input(shape=(28, 28, 1), name="img")
x = tf.keras.layers.Conv2D(16, 3, activation="relu")(encoder_input)
x = tf.keras.layers.Conv2D(32, 3, activation="relu")(x)
x = tf.keras.layers.MaxPooling2D(3)(x)
x = tf.keras.layers.Conv2D(32, 3, activation="relu")(x)
x = tf.keras.layers.Lambda(lambda x: tf.einsum('...ijk->...jik', x), name='Einsum')(x)  # tensorflow op, just swap height and width for no reason
x = tf.keras.layers.Conv2D(16, 3, activation="relu")(x)
model = tf.keras.Model(encoder_input, x)
print(model.summary())

Model: "model_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 img (InputLayer)            [(None, 28, 28, 1)]       0         
                                                                 
 conv2d_4 (Conv2D)           (None, 26, 26, 16)        160       
                                                                 
 conv2d_5 (Conv2D)           (None, 24, 24, 32)        4640      
                                                                 
 max_pooling2d_1 (MaxPooling  (None, 8, 8, 32)         0         
 2D)                                                             
                                                                 
 conv2d_6 (Conv2D)           (None, 6, 6, 32)          9248      
                                                                 
 Einsum (Lambda)             (None, 6, 6, 32)          0         
                                                                 
 conv2d_7 (Conv2D)           (None, 4, 4, 16)          4624      
                                                                 
=================================================================
Total params: 18,672
Trainable params: 18,672
Non-trainable params: 0
_________________________________________________________________
None

例如您使用tf.Reduce_sum在axis = 0中，您将丢失批处理维度（none），这对于keras < /代码>型号：

import tensorflow as tf

encoder_input = tf.keras.Input(shape=(28, 28, 1), name="img")
x = tf.keras.layers.Conv2D(16, 3, activation="relu")(encoder_input)
x = tf.reduce_sum(x, axis=0)
model = tf.keras.Model(encoder_input, x)
print(model.summary())

Model: "model_4"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 img (InputLayer)            [(None, 28, 28, 1)]       0         
                                                                 
 conv2d_12 (Conv2D)          (None, 26, 26, 16)        160       
                                                                 
 tf.math.reduce_sum_1 (TFOpL  (26, 26, 16)             0         
 ambda)                                                          
                                                                 
=================================================================
Total params: 160
Trainable params: 160
Non-trainable params: 0
_________________________________________________________________
None