Can you mix keras.layers output with tensorflow operations?

Question

I know that output of keras layers (like keras.layers.Dense()) produce so-called 'keras tensors'. Also, there are 'tensorflow tensors' that are produced by tensorflow operations (like tf.reduce_sum()).

Some functionality can be delivered by only one of those approaches, so it is obvious that I will have to mix them sometimes to do the calculation. In my opinion, mixing keras layers code with tf ops looks far from beatiful but thats another topic.

My question is - is it doable to mix keras layers and tensorflow ops? And will there be any problems related to mixing keras and tensorflow tensors that they produce?

---

Lets consider an example of custom class, don't look deep into the calculation, it makes no sense. The point is to show what I'm talking about.

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

---

Another example, this time not inside custom layer:

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)

---

I know that there is such thing as keras.layers.Lambda layer, but sometimes I need to use like 90% of tensorflow ops and only 10% of calculations done by keras layers (because there is no tf op alternative and I don't want to implement my own every time). In that case using lambda layer makes little sense.

Is there a good approach for writing complex models (where implementation lays beyond just stacking keras existing layers) in tensorflow 2.X?

Answer 1

IMO, it actually does not matter how you mix Tensorflow operations with Keras layers as long as you preserve the batch dimension and generally the tensor shape. You might want to, for example, wrap the tf operations inside Lambda layers to set some meta-data like the names, but that depends on your taste:

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

For example, if you use tf.reduce_sum with axis=0, you lose the batch dimension (None), which is problematic for a Keras model:

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