Fixing LocallyConnected2D and LocallyConnected1D layer Save Model to tf issue

[ashutosh1919]

Fixes #48584 and #47689 .

My understanding:

The model save is failing because at the time of saving, call() of LocallyConnected2D will be invoked which is not passed with any inputs (None). That is why compute_output_shape and local_conv_matmul fails and thus save fails.

Solution that I propose:

We particularly don't require any input data at the time of saving. So, I am saving the input_shape at the time when LocallyConnected2D instance is build() and whenever the parameter inputs is None in call(), I am replacing inputs with the dummy tensor with shape input_shape which was saved in build().

Any suggestions on optimisation of solution are welcomed.

cc @mihaimaruseac , @bhack , @ AdityaKane2001

[bhack]

Can you add a test to cover the mentioned tickets case?

[AdityaKane2001]

@ashutosh1919
I ran your code in colab. I made a model with one LocallyConnected2D layer. I passed a random numpy array to the model. Then I tried to save the model, but the notebook crashed due to excessive RAM usage. Gist here.

[ashutosh1919]

@AdityaKane2001 , how are you running code in my branch? I see direct import import tensorflow as tf in the gist you shared above.
Also, try reducing the size of the np array. I tried with (10,10,10,3) and works fine in your gist.

[AdityaKane2001]

@ashutosh1919
It's sort of a small hack but it works for small changes. I replaced the changed files in local.py of the colab session with your code before importing tensorflow. It then uses the new code.
Actually, it's my fault, I forgot to add implementation = 2 for LocallyConnected2D .

[AdityaKane2001]

Updated gist here. I have used your code as mentioned above. Still facing the same error.

[ashutosh1919]

@AdityaKane2001 , I have resolved the error and added test cases. And btw, I think you can't just change the code for local.py in collab. I think you need to bazel build again if you change the code.

@bhack , I have added the test case to save all 3 implementation of models and the problem is there for LocallyConnected1D as well. So, I have fixed the code in that too. Bazel tests are passing in my local. Can you please approve the PR so that the external build tests can run?

Solution (Attempt 2):

Actually at the time of save(), it is not that No inputs are passed, Inputs are passed but the shape of that is something like (None, None, None, 2) which is not the conventional shape (None, 4, 4, 2). So, the solution is,

• To store the shape of input whenever the layer is built.
• Whenever call() is invoked, check whether any of the dimension shape other than first one is None. If yes, then replace input_shape with stored_input_shape otherwise proceed as it is.

Any other suggestions on this are welcome. Also I would suggest you to please review code and suggest changes because I have raised ValueErrors as well as tf_logging.warn and I am not sure about the messages that I am populating in that. It may need correction.

cc @mihaimaruseac @fchollet

[AdityaKane2001]

@ashutosh1919
AFAIK, it works because we are altering the python code, which is in the very last layer of the API system.

Anyways, apart from that, everything looks good. 👍

[k-w-w]

Thanks for the fix!

It looks like the LocallyConnected layer is failing as you noted, the input shapes have None in them. This is due to the input_spec being used as the inputs (for most Keras built-in layers using the input spec is fine, but not for LocallyConnected layers):

A simple fix would be overwriting this property to the LocallyConnected1D and LocallyConnected2D:

class LocallyConnected1D:
  ...
  @property
  def _use_input_spec_as_call_signature(self):
    return False

[ashutosh1919]

@k-w-w , Thanks for straight forward and easy solution. I got to learn something new here.
I have made required changes. Please review.

[k-w-w]

Thanks !