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

Question

我知道 keras 层（如 keras.layers.Dense()）的输出会产生所谓的 'keras tensors'。此外，还有 'tensorflow tensors' 是由 tensorflow 操作产生的（如 tf.reduce_sum()）。

某些功能只能通过这些方法中的一种来提供，因此很明显，我有时不得不混合使用它们 来进行计算。在我看来，将 keras 层代码与 tf ops 混合看起来远非美丽，但那是另一个话题。

我的问题是 - 是否可以混合使用 keras 层和 tensorflow 操作？并且他们产生的混合keras和tensorflow张量会不会有任何问题？

让我们考虑一个自定义的例子class，不要深入研究计算，它没有意义。重点是展示我在说什么。

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 层这样的东西，但有时我需要使用像 90% 这样的 tensorflow ops 并且只需要 10 % 由 keras 层完成的计算（因为没有 tf op 替代方案，我不想每次都实现我自己的）。在那种情况下，使用 lambda 层毫无意义。

在 tensorflow 2.X 中是否有编写复杂模型（其中实现不仅仅是堆叠 keras 现有层）的好方法？

Answer 1

IMO，实际上，将 Tensorflow 操作与 Keras 层混合使用并不重要，只要保留批次维度和通常的张量形状即可。例如，您可能希望将 tf 操作包装在 Lambda 层中以设置一些 meta-data 名称，但这取决于您的喜好：

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

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

Can you mix keras.layers output with tensorflow operations?

python

tensorflow

keras

deep-learning

tensorflow2.0