Incorporating a TensorHub model into my Keras model

user_17429148 0
Viewed 114

TensorFlow 2.7, Keras 2.7

I am trying to use an existing TFHub model as a layer in my model. Wrapped it by a custom keras layer but probably missed something around the batch size. Wrote a simple version of it below. The model below receives [ BATCH_SIZE, 224, 224, 3 ], uses a TFHub model to generate one representation, another simple layer to generate another representation. Then it concatenates both, uses a Dense layer and outputs [BATCH_SIZE, 10].

It works fine with batch size = 1 but with batch size > 1 predict and evaluate works but fit returns an error related to the last Dense layer receiving an incorrect size input.

Simple model and call code:

def simple_model( ):

    input = Input(shape=(224, 224, 3))

    representation = Conv2D(1,(224,224))(input)

    # Prepare input for depth estimation model
    resized_for_midas = tf.image.resize(input, (384, 384))
    transposed = tf.transpose(resized_for_midas, [0, 3, 1, 2])

    depth_estimation = tfhub.KerasLayer('https://tfhub.dev/intel/midas/v2/2', signature = 'serving_default',
    tags=['serve'])(transposed)
    depth_estimation_reshaped = tf.expand_dims(depth_estimation, axis=-1) # Adding 1 in the end so the shape will be [batch_size, 384, 384, 1]
    depth_estimation_repmat = Conv2D(3,(3,3),padding='same')(depth_estimation_reshaped) # Repeat the depth estimation single channel 3 times to match ResNet input using a convolution layer
    depth_estimation_resized = tf.image.resize(depth_estimation_repmat, (224, 224))

    depth_estimation_representation = Conv2D(1,(224,224))(depth_estimation_resized)

    # Concatenate representations
    representation_full = tf.concat([representation, depth_estimation_representation], axis=-1)
    
    flat = tf.reshape(representation_full, (-1, representation_full.shape[-1]))

    # Outputs
    output = Dense(10, input_shape=(-1, flat.shape[-1]),
        activation='linear')(flat)

    model = Model(inputs=input, outputs=output)
    
    return model

model = simple_model()
model.compile(loss='mae')
batch_size = 1
model(np.random.rand(batch_size, 224, 224, 3)) # Works
model.evaluate(np.random.rand(batch_size, 224, 224, 3), np.random.rand(batch_size, 10)) # Works
model.fit(np.random.rand(batch_size, 224, 224, 3), np.random.rand(batch_size, 10)) # Works
batch_size = 2
model(np.random.rand(batch_size, 224, 224, 3)) # Works
model.evaluate(np.random.rand(batch_size, 224, 224, 3), np.random.rand(batch_size, 10)) # Works
model.fit(np.random.rand(batch_size, 224, 224, 3), np.random.rand(batch_size, 10)) # Fails

Error:

Traceback (most recent call last): File "", line 1, in File "/home/dani/projects/venv/lib/python3.8/site-packages/wandb/integration/keras/keras.py", line 150, in new_v2 return old_v2(*args, **kwargs) File "/home/dani/projects/venv/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler raise e.with_traceback(filtered_tb) from None File "/home/dani/projects/venv/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 58, in quick_execute tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name, tensorflow.python.framework.errors_impl.InvalidArgumentError: Matrix size-incompatible: In[0]: [2,2], In[1]: [1,10] [[node gradient_tape/model_1/dense_1/MatMul/MatMul_1 (defined at /home/dani/projects/venv/lib/python3.8/site-packages/keras/optimizer_v2/optimizer_v2.py:464) ]] [Op:__inference_train_function_30163]

Errors may have originated from an input operation. Input Source operations connected to node gradient_tape/model_1/dense_1/MatMul/MatMul_1: In[0] model_1/tf.reshape_1/Reshape (defined at /home/dani/projects/venv/lib/python3.8/site-packages/keras/layers/core/tf_op_layer.py:261)
In[1] gradient_tape/mean_absolute_error/sub/Reshape:

keras tensorflow python
1 Answers

@Dani This is because the tfhub.KerasLayer you are using in your code is built for a single image.

It is clearly mentioned in that tfhub page https://tfhub.dev/intel/midas/v2/2

Convolutional neural network for monocular depth estimation from a single RGB image.

Also, check the model.summary which shows the discrepancy

Model: "model"
__________________________________________________________________________________________________
 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 input_1 (InputLayer)           [(None, 224, 224, 3  0           []                               
                                )]                                                                
                                                                                                  
 tf.image.resize (TFOpLambda)   (None, 384, 384, 3)  0           ['input_1[0][0]']                
                                                                                                  
 tf.compat.v1.transpose (TFOpLa  (None, 3, 384, 384)  0          ['tf.image.resize[0][0]']        
 mbda)                                                                                            
                                                                                                  
 keras_layer (KerasLayer)       (1, 384, 384)        0           ['tf.compat.v1.transpose[0][0]'] 
                                                                                                  
 tf.expand_dims (TFOpLambda)    (1, 384, 384, 1)     0           ['keras_layer[0][0]']            
                                                                                                  
 conv2d_1 (Conv2D)              (1, 384, 384, 3)     30          ['tf.expand_dims[0][0]']         
                                                                                                  
 tf.image.resize_1 (TFOpLambda)  (1, 224, 224, 3)    0           ['conv2d_1[0][0]']               
                                                                                                  
 conv2d (Conv2D)                (None, 1, 1, 1)      150529      ['input_1[0][0]']                
                                                                                                  
 conv2d_2 (Conv2D)              (1, 1, 1, 1)         150529      ['tf.image.resize_1[0][0]']      
                                                                                                  
 tf.concat (TFOpLambda)         (1, 1, 1, 2)         0           ['conv2d[0][0]',                 
                                                                  'conv2d_2[0][0]']               
                                                                                                  
 tf.reshape (TFOpLambda)        (1, 2)               0           ['tf.concat[0][0]']              
                                                                                                  
 dense (Dense)                  (1, 10)              30          ['tf.reshape[0][0]']             
                                                                                                  
==================================================================================================
Total params: 301,118
Trainable params: 301,118
Non-trainable params: 0
__________________________________________________________________________________________________

Please check the gist here

I guess you need to find a layer that supports multiple images as a batch. Thanks!

user_142906810
Related