search

keras-team / keras

Deep Learning for humans
http://keras.io/
Apache License 2.0
61.35k stars 19.39k forks source link

AssertionError when Saving Model as SavedModel Format with image augmentation layers #19100

Open lxzheng opened 5 months ago

lxzheng commented 5 months ago

Environment:

Issue Description:

Encountered an AssertionError when saving a Keras model as SavedModel format, specifically when the model includes image augmentation layers like RandomFlip or RandomRotation. 

Steps to Reproduce:

The following minimal code example demonstrates this error. The error occurs when executing the tf.saved_model.save function.

import tensorflow as tf
import keras

(train_images, train_labels), (test_images, test_labels) = keras.datasets.fashion_mnist.load_data()

test_model = keras.Sequential([
    keras.layers.RandomFlip("horizontal_and_vertical"),  
    keras.layers.RandomRotation(0.1),  

    keras.layers.Flatten(),
    keras.layers.Dense(64, activation='relu'),

    keras.layers.Dense(10, activation='softmax')
])
test_model.compile(optimizer='adam',
              loss=keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])
test_model.fit(train_images, train_labels, epochs=1)
tf.saved_model.save(test_model, "test")

Error Message:

AssertionError Traceback (most recent call last) Cell In[264], line 1 ----> 1 tf.saved_model.save(test_model, "test")

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/saved_model/save.py:1336, in save(obj, export_dir, signatures, options) 1334 # pylint: enable=line-too-long 1335 metrics.IncrementWriteApi(_SAVE_V2_LABEL) -> 1336 save_and_return_nodes(obj, export_dir, signatures, options) 1338 metrics.IncrementWrite(write_version="2")

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/saved_model/save.py:1371, in save_and_return_nodes(obj, export_dir, signatures, options, experimental_skip_checkpoint) 1367 saved_model = saved_model_pb2.SavedModel() 1368 meta_graph_def = saved_model.metagraphs.add() 1370 , exported_graph, object_saver, asset_info, saved_nodes, node_paths = ( -> 1371 _build_meta_graph(obj, signatures, options, meta_graph_def)) 1372 saved_model.saved_model_schema_version = ( 1373 constants.SAVED_MODEL_SCHEMA_VERSION) 1375 # Write the checkpoint, copy assets into the assets directory, and write out 1376 # the SavedModel proto itself.

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/saved_model/save.py:1584, in _build_meta_graph(obj, signatures, options, meta_graph_def) 1557 """Creates a MetaGraph under a save context. 1558 1559 Args: (...) 1580 saveable_view.node_paths: _SaveableView paths. 1581 """ 1583 with save_context.save_context(options): -> 1584 return _build_meta_graph_impl(obj, signatures, options, meta_graph_def)

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/saved_model/save.py:1509, in _build_meta_graph_impl(obj, signatures, options, meta_graph_def) 1507 saveable_view = _SaveableView(augmented_graph_view, options) 1508 object_saver = checkpoint.TrackableSaver(augmented_graph_view) -> 1509 asset_info, exported_graph = _fill_meta_graph_def( 1510 meta_graph_def=meta_graph_def, 1511 saveable_view=saveable_view, 1512 signature_functions=signatures, 1513 namespace_whitelist=options.namespace_whitelist, 1514 save_custom_gradients=options.experimental_custom_gradients, 1515 create_saver=not options.experimental_skip_saver, 1516 defaults=defaults, 1517 ) 1518 if options.function_aliases: 1519 function_aliases = meta_graph_def.meta_info_def.function_aliases

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/saved_model/save.py:879, in _fill_meta_graph_def(meta_graph_def, saveable_view, signature_functions, namespace_whitelist, save_custom_gradients, create_saver, defaults) 877 with exported_graph.as_default(): 878 object_map, tensor_map, asset_info = saveable_view.map_resources() --> 879 signatures = _generate_signatures(signature_functions, object_map, defaults) 880 if save_custom_gradients: 881 # Custom gradients functions must be traced in the same context as the 882 # when they are registered. 883 _trace_gradient_functions(exported_graph, saveable_view)

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/saved_model/save.py:654, in _generate_signatures(signature_functions, object_map, defaults) 646 mapped_inputs, exterior_argument_placeholders = ( 647 _map_function_arguments_to_created_inputs( 648 argument_inputs, signature_key, function.name, defaults 649 ) 650 ) 651 kwarg_names = list( 652 sorted( 653 object_map[function].function.structured_input_signature[1].keys())) --> 654 outputs = object_map[function](**{ 655 kwarg_name: mapped_input 656 for kwarg_name, mapped_input in zip(kwarg_names, mapped_inputs) 657 }) 658 signatures[signature_key] = signature_def_utils.build_signature_def( 659 _tensor_dict_to_tensorinfo(exterior_argument_placeholders), 660 _tensor_dict_to_tensorinfo(outputs), 661 method_name=signature_constants.PREDICT_METHOD_NAME, 662 defaults=defaults.get(signature_key, None), 663 ) 664 return signatures

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py:45, in ExportedConcreteFunction.call(self, *args, **kwargs) 39 bound_arguments = function_type_utils.canonicalize_function_inputs( 40 args, kwargs, self.function._function_type 41 ) 42 filtered_flat_args = self.function._function_type.unpack_inputs( 43 bound_arguments 44 ) ---> 45 export_captures = _map_captures_to_created_tensors( 46 self.function.graph.captures, self.tensor_map, self.function) 47 return self.function._call_flat(filtered_flat_args, export_captures)

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py:74, in _map_captures_to_created_tensors(original_captures, tensor_map, function) 72 mapped_resource = tensor_map.get(exterior, None) 73 if mapped_resource is None: ---> 74 _raise_untracked_capture_error(function.name, exterior, interior) 75 export_captures.append(mapped_resource) 76 return export_captures

File /usr/local/lib/python3.11/dist-packages/tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py:98, in _raise_untracked_capture_error(function_name, capture, internal_capture, node_path) 96 if internal_capture is not None: 97 msg += f"\n\tInternal Tensor = {internal_capture}" ---> 98 raise AssertionError(msg)

AssertionError: Tried to export a function which references an 'untracked' resource. TensorFlow objects (e.g. tf.Variable) captured by functions must be 'tracked' by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly. See the information below: Function name = b'__inference_signature_wrapper_serving_default_2627467' Captured Tensor = <ResourceHandle(name="seed_generator_state/4022", device="/job:localhost/replica:0/task:0/device:CPU:0", container="Anonymous", type="tensorflow::Var", dtype and shapes : "[ DType enum: 22, Shape: [2] ]")> Trackable referencing this tensor = <tf.Variable 'seed_generator_state:0' shape=(2,) dtype=uint32> Internal Tensor = Tensor("2627459:0", shape=(), dtype=resource)

SuryanarayanaY commented 5 months ago

Hi @lxzheng ,

I have replicated the reported behaviour with augmentation layers. Without augmentation layers it works fine. Attached gist for reference.

lxzheng commented 5 months ago

@SuryanarayanaY 

Thank you for confirming the behavior with the augmentation layers and for providing the reference gist. Of course, models without augmentation layers can indeed be saved successfully in the SavedModel format. However, I have discovered an additional issue: even a Keras model without data augmentation layers encounters an error when converting to TensorFlow Lite  format.

I have opened a new GitHub issue (#19108) to report this bug, which occurs during the TFLite conversion process regardless of the presence of augmentation layers.