Using 3d sparse input with partial shape in Keras/Tensorflow gives error

Question

I am trying to create a dense neural network where my input is a sparse 3d matrix. When converted to a dense matrix the shape is (2, None, n) (where n is a number of features and is fixed.) My keras architecture works fine when I am using dense inputs. However, to save memory I am trying to use sparse tensors as input.

Here is my code

import numpy as np
import tensorflow as tf

input = tf.keras.Input(batch_shape=(2,None,5))
x = tf.keras.layers.Dropout(0.1)(input)
x = tf.keras.layers.Dense(1)(x)

model = tf.keras.models.Model(inputs=[input], outputs=[x])
model.compile(loss='mse')
print(model.summary())

dummy_input = np.random.random((2,10,5))
dummy_sp = tf.sparse.from_dense(dummy_input)

dummy_output= np.random.random((2,10,1))

model.fit(x=dummy_sp, y=dummy_output, epochs=1)

The above code works fine when I use the x=dummy_input. Howeverm when I switch to the sparse inputs dummy_sp then I get the following error

Model: "model"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_1 (InputLayer)        [(2, None, 5)]            0

 dropout (Dropout)           (2, None, 5)              0

 dense (Dense)               (2, None, 1)              6

=================================================================
Total params: 6
Trainable params: 6
Non-trainable params: 0
_________________________________________________________________
None
Traceback (most recent call last):
  File "c:\Users\099391\OneDrive\Documents\Projects\NP_E-QSI\srcc\test.py", line 17, in <module>
    model.fit(x=dummy_sp, y=dummy_output, epochs=1)
  File "C:\Users\099391\Anaconda3\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\099391\Anaconda3\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1147, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:

    File "C:\Users\099391\Anaconda3\lib\site-packages\keras\engine\training.py", line 1021, in train_function  *
        return step_function(self, iterator)
    File "C:\Users\099391\Anaconda3\lib\site-packages\keras\engine\training.py", line 1010, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\099391\Anaconda3\lib\site-packages\keras\engine\training.py", line 1000, in run_step  **
        outputs = model.train_step(data)
    File "C:\Users\099391\Anaconda3\lib\site-packages\keras\engine\training.py", line 859, in train_step
        y_pred = self(x, training=True)
    File "C:\Users\099391\Anaconda3\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
        raise e.with_traceback(filtered_tb) from None

    TypeError: Exception encountered when calling layer "dropout" (type Dropout).

    Failed to convert elements of SparseTensor(indices=Tensor("DeserializeSparse:0", shape=(None, 3), dtype=int64), values=Tensor("model/Cast:0", shape=(None,), dtype=float32), dense_shape=Tensor("stack:0", shape=(3,), dtype=int64)) to Tensor. Consider casting elements to a supported type. See https://www.tensorflow.org/api_docs/python/tf/dtypes for supported TF dtypes.

    Call arguments received:
      • inputs=<tensorflow.python.framework.sparse_tensor.SparseTensor object at 0x000002423CD78E20>
      • training=True

Answer 1

Based on the tf-documentation, you can still pass the sparse tensor even if the sparse argument is set as False, as long as the input data is a sparse matrix.

sparse : A boolean specifying whether the placeholder to be created is sparse. Only one of 'ragged' and 'sparse' can be True. Note that, if sparse is False, sparse tensors can still be passed into the input - they will be densified with a default value of 0.

input = tf.keras.Input(batch_shape=(2,None,5), sparse=False)
x = tf.keras.layers.Dropout(0.1)(input)
x = tf.keras.layers.Dense(1)(x)
model = tf.keras.models.Model(inputs=[input], outputs=[x])
print(model.summary())
# OK 

Update

Based on the comment below.

class toDense(keras.layers.Layer):
    def call(self, input):
        if isinstance(input, tf.sparse.SparseTensor):
            return tf.sparse.to_dense(input)
        return input

Place it right after Input and before the Dropout layer of the above model. Like,

input = tf.keras.Input(batch_shape=(2,None,5))
x = toDense()(input)
x = tf.keras.layers.Dropout(0.1)(x)
...

Next, you can do

dummy_input = np.random.random((2,10,5))
dummy_sp = tf.sparse.from_dense(dummy_input)
dummy_output= np.random.random((2,10,1))

model(dummy_sp).shape 
>> TensorShape([2, 10, 1])

model.compile(loss='mse')
model.fit(x=dummy_sp, y=dummy_output, epochs=1)
1/1 [=====] - 0s 416ms/step - loss: 0.2746