search

tensorflow / neural-structured-learning

Training neural models with structured signals.
https://www.tensorflow.org/neural_structured_learning
Apache License 2.0
980 stars 189 forks source link

StagingError appears for TensorFlow 2.2.0-rc3 #57

Closed sayakpaul closed 4 years ago

sayakpaul commented 4 years ago

Hi

I am primarily trying to train an adversarial variant of a custom image classifier. The code works perfectly fine when I am on TensorFlow 2.1.0, 2.2.0-rc2. But the same code breaks on TensorFlow 2.2.0-rc3 producing:

StagingError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:571 train_function  *
        outputs = self.distribute_strategy.run(
    /usr/local/lib/python3.6/dist-packages/neural_structured_learning/keras/adversarial_regularization.py:627 call  *
        outputs, labeled_loss, metrics, tape = self._forward_pass(
    /usr/local/lib/python3.6/dist-packages/neural_structured_learning/keras/adversarial_regularization.py:606 _forward_pass  *
        outputs = self.base_model(inputs, **base_model_kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py:927 __call__  **
        outputs = call_fn(cast_inputs, *args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py:719 call
        convert_kwargs_to_constants=base_layer_utils.call_context().saving)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py:826 _run_internal_graph
        inputs = self._flatten_to_reference_inputs(inputs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py:926 _flatten_to_reference_inputs
        return [tensors[inp._keras_history.layer.name] for inp in ref_inputs]
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py:926 <listcomp>
        return [tensors[inp._keras_history.layer.name] for inp in ref_inputs]

    KeyError: 'input_1'

Here's the Colab Gist that produces this issue. Here's the Colab Notebook that shows that the code works for TensorFlow 2.2.0-rc2.

csferng commented 4 years ago

Thanks for bringing this up. In short, please change the dictionary key of input images from "image" to EXTRACTOR.input_names[0] or base_adv_model.input_names[0]. The reason is as follows:

Recently TensorFlow made a change on handling dictionary-format input when a model is called directly (e.g. model(input_dict)). Before the change, the input_dict is flattened to a list (ordered by keys) and fed to tensors in model.inputs. After the change, the input_dict is looked up based on model.input_names. For tf.keras.applications.MobileNetV2, its input_names is generated automatically and different from our hand-chosen name like "image". The mismatch is fine in previous versions of TensorFlow because the names are ignored, but is causing KeyError as you saw in the latest version of TensorFlow. To overcome this issue, please use the generated input name as the dictionary key.

sayakpaul commented 4 years ago

I see. I will try it out and will get back but that may not be timely, apologies for that.