Question: I am getting an error when trying to save a model with data augmentation layers with Tensorflow version 2.7.0. Here is the code of data augmentation: input_shape_rgb = (img_height, img_width, 3) data_augmentation_rgb = tf.keras.Sequential( [ layers.RandomFlip("horizontal"), layers.RandomFlip("vertical"), layers.RandomRotation(0.5), layers.RandomZoom(0.5), layers.RandomContrast(0.5), RandomColorDistortion(name='random_contrast_brightness/none'), ] )  Now I build my model like this: # Build the model input_shape = (img_height, img_width, 3) model = Sequential([ layers.Input(input_shape), data_augmentation_rgb, layers.Rescaling((1./255)), layers.Conv2D(16, kernel_size, padding=padding, activation='relu', strides=1, data_format='channels_last'), layers.MaxPooling2D(), layers.BatchNormalization(), layers.Conv2D(32, kernel_size, padding=padding, activation='relu'), # best 4 layers.MaxPooling2D(), layers.BatchNormalization(), layers.Conv2D(64, kernel_size, padding=padding, activation='relu'), # best 3 layers.MaxPooling2D(), layers.BatchNormalization(), layers.Conv2D(128, kernel_size, padding=padding, activation='relu'), # best 3 layers.MaxPooling2D(), layers.BatchNormalization(), layers.Flatten(), layers.Dense(128, activation='relu'), # best 1 layers.Dropout(0.1), layers.Dense(128, activation='relu'), # best 1 layers.Dropout(0.1), layers.Dense(64, activation='relu'), # best 1 layers.Dropout(0.1), layers.Dense(num_classes, activation = 'softmax') ]) model.compile(loss='categorical_crossentropy', optimizer='adam',metrics=metrics) model.summary()  Then after the training is done I just make: model.save("./")  And I'm getting this error: --------------------------------------------------------------------------- KeyError Traceback (most recent call last) <ipython-input-84-87d3f09f8bee> in <module>() ----> 1 model.save("./") /usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs) 65 except Exception as e: # pylint: disable=broad-except 66 filtered_tb = _process_traceback_frames(e.__traceback__) ---> 67 raise e.with_traceback(filtered_tb) from None 68 finally: 69 del filtered_tb /usr/local/lib/python3.7/dist- packages/tensorflow/python/saved_model/function_serialization.py in serialize_concrete_function(concrete_function, node_ids, coder) 66 except KeyError: 67 raise KeyError( ---> 68 f"Failed to add concrete function '{concrete_function.name}' to object-" 69 f"based SavedModel as it captures tensor {capture!r} which is unsupported" 70 " or not reachable from root. " KeyError: "Failed to add concrete function 'b'__inference_sequential_46_layer_call_fn_662953'' to object-based SavedModel as it captures tensor <tf.Tensor: shape=(), dtype=resource, value=<Resource Tensor>> which is unsupported or not reachable from root. One reason could be that a stateful object or a variable that the function depends on is not assigned to an attribute of the serialized trackable object (see SaveTest.test_captures_unreachable_variable)."  I inspected the reason of getting this error by changing the architecture of my model and I just found that reason came from the data_augmentation layer since the RandomFlip and RandomRotation and others are changed from layers.experimental.prepocessing.RandomFlip to layers.RandomFlip, but still the error appears.

Answer: This seems to be a bug in Tensorflow 2.7 when using model.save combined with the parameter save_format="tf", which is set by default. The layers RandomFlip, RandomRotation, RandomZoom, and RandomContrast are causing the problems, since they are not serializable. Interestingly, the Rescaling layer can be saved without any problems. A workaround would be to simply save your model with the older Keras H5 format model.save("test", save_format='h5'): import tensorflow as tf import numpy as np class RandomColorDistortion(tf.keras.layers.Layer): def __init__(self, contrast_range=[0.5, 1.5], brightness_delta=[-0.2, 0.2], **kwargs): super(RandomColorDistortion, self).__init__(**kwargs) self.contrast_range = contrast_range self.brightness_delta = brightness_delta def call(self, images, training=None): if not training: return images contrast = np.random.uniform( self.contrast_range[0], self.contrast_range[1]) brightness = np.random.uniform( self.brightness_delta[0], self.brightness_delta[1]) images = tf.image.adjust_contrast(images, contrast) images = tf.image.adjust_brightness(images, brightness) images = tf.clip_by_value(images, 0, 1) return images def get_config(self): config = super(RandomColorDistortion, self).get_config() config.update({"contrast_range": self.contrast_range, "brightness_delta": self.brightness_delta}) return config input_shape_rgb = (256, 256, 3) data_augmentation_rgb = tf.keras.Sequential( [ tf.keras.layers.RandomFlip("horizontal"), tf.keras.layers.RandomFlip("vertical"), tf.keras.layers.RandomRotation(0.5), tf.keras.layers.RandomZoom(0.5), tf.keras.layers.RandomContrast(0.5), RandomColorDistortion(name='random_contrast_brightness/none'), ] ) input_shape = (256, 256, 3) padding = 'same' kernel_size = 3 model = tf.keras.Sequential([ tf.keras.layers.Input(input_shape), data_augmentation_rgb, tf.keras.layers.Rescaling((1./255)), tf.keras.layers.Conv2D(16, kernel_size, padding=padding, activation='relu', strides=1, data_format='channels_last'), tf.keras.layers.MaxPooling2D(), tf.keras.layers.BatchNormalization(), tf.keras.layers.Conv2D(32, kernel_size, padding=padding, activation='relu'), # best 4 tf.keras.layers.MaxPooling2D(), tf.keras.layers.BatchNormalization(), tf.keras.layers.Conv2D(64, kernel_size, padding=padding, activation='relu'), # best 3 tf.keras.layers.MaxPooling2D(), tf.keras.layers.BatchNormalization(), tf.keras.layers.Conv2D(128, kernel_size, padding=padding, activation='relu'), # best 3 tf.keras.layers.MaxPooling2D(), tf.keras.layers.BatchNormalization(), tf.keras.layers.Flatten(), tf.keras.layers.Dense(128, activation='relu'), # best 1 tf.keras.layers.Dropout(0.1), tf.keras.layers.Dense(128, activation='relu'), # best 1 tf.keras.layers.Dropout(0.1), tf.keras.layers.Dense(64, activation='relu'), # best 1 tf.keras.layers.Dropout(0.1), tf.keras.layers.Dense(5, activation = 'softmax') ]) model.compile(loss='categorical_crossentropy', optimizer='adam') model.summary() model.save("test", save_format='h5')  Loading your model with your custom layer would look like this then: model = tf.keras.models.load_model('test.h5', custom_objects={'RandomColorDistortion': RandomColorDistortion})  where RandomColorDistortion is the name of your custom layer.