(AttributeError: 'NoneType' object has no attribute 'get' ) while loading saved keras model with .h5 extension in tensorflow 2.1

Question

I have a keras model making use of feature_column api of tensorflow , I am able to save the model in .h5 extension but gets following error in colab while loading the saved model.

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-38-8a2d51a054f6> in <module>()
----> 1 new_model = tf.keras.models.load_model('model.h5')

13 frames

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    318       obj = _GLOBAL_CUSTOM_OBJECTS[object_name]
    319     else:
--> 320       obj = module_objects.get(object_name)
    321       if obj is None:
    322         raise ValueError('Unknown ' + printable_module_name + ':' + object_name)

AttributeError: 'NoneType' object has no attribute 'get'

following is my code and i am using tensorlow 2.1

import pandas as pd
import tensorflow as tf 
import numpy as np

feature_layer = tf.keras.layers.DenseFeatures(feature_columns)

def model_builder():


      tf.keras.backend.clear_session()
      model = tf.keras.Sequential()
      model.add(feature_layer)
      model.add(tf.keras.layers.Dense(32, activation='relu'))
      model.add(tf.keras.layers.Dense(32, activation='relu'))
      model.add(tf.keras.layers.Dense(1))


      model.compile(optimizer = tf.keras.optimizers.RMSprop(0.001),
                    loss=tf.keras.losses.Huber(),
                    metrics=['mae'])


      return model

model = model_builder()
model.fit(train_ds , epochs=5,verbose=1)

model.save('model.h5')

new_model = tf.keras.models.load_model('model.h5')

I get the following error while running code in local environment ie in conda environment using jupyter notebook. But local env is having tf 2.0

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-23-fbbaf2c35e78> in <module>
----> 1 new_model = tf.keras.models.load_model('model1.h5')

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/save.py in load_model(filepath, custom_objects, compile)
    144   if (h5py is not None and (
    145       isinstance(filepath, h5py.File) or h5py.is_hdf5(filepath))):
--> 146     return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
    147 
    148   if isinstance(filepath, six.string_types):

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/hdf5_format.py in load_model_from_hdf5(filepath, custom_objects, compile)
    166     model_config = json.loads(model_config.decode('utf-8'))
    167     model = model_config_lib.model_from_config(model_config,
--> 168                                                custom_objects=custom_objects)
    169 
    170     # set weights

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/saving/model_config.py in model_from_config(config, custom_objects)
     53                     '`Sequential.from_config(config)`?')
     54   from tensorflow.python.keras.layers import deserialize  # pylint: disable=g-import-not-at-top
---> 55   return deserialize(config, custom_objects=custom_objects)
     56 
     57 

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/serialization.py in deserialize(config, custom_objects)
    100       module_objects=globs,
    101       custom_objects=custom_objects,
--> 102       printable_module_name='layer')

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    189             custom_objects=dict(
    190                 list(_GLOBAL_CUSTOM_OBJECTS.items()) +
--> 191                 list(custom_objects.items())))
    192       with CustomObjectScope(custom_objects):
    193         return cls.from_config(cls_config)

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/sequential.py in from_config(cls, config, custom_objects)
    367     for layer_config in layer_configs:
    368       layer = layer_module.deserialize(layer_config,
--> 369                                        custom_objects=custom_objects)
    370       model.add(layer)
    371     if not model.inputs and build_input_shape:

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/serialization.py in deserialize(config, custom_objects)
    100       module_objects=globs,
    101       custom_objects=custom_objects,
--> 102       printable_module_name='layer')

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    189             custom_objects=dict(
    190                 list(_GLOBAL_CUSTOM_OBJECTS.items()) +
--> 191                 list(custom_objects.items())))
    192       with CustomObjectScope(custom_objects):
    193         return cls.from_config(cls_config)

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/feature_column/feature_column_v2.py in from_config(cls, config, custom_objects)
    450     config_cp = config.copy()
    451     config_cp['feature_columns'] = serialization.deserialize_feature_columns(
--> 452         config['feature_columns'], custom_objects=custom_objects)
    453 
    454     return cls(**config_cp)

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/feature_column/serialization.py in deserialize_feature_columns(configs, custom_objects)
    188   return [
    189       deserialize_feature_column(c, custom_objects, columns_by_name)
--> 190       for c in configs
    191   ]

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/feature_column/serialization.py in <listcomp>(.0)
    188   return [
    189       deserialize_feature_column(c, custom_objects, columns_by_name)
--> 190       for c in configs
    191   ]

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/feature_column/serialization.py in deserialize_feature_column(config, custom_objects, columns_by_name)
    141       cls_config,
    142       custom_objects=custom_objects,
--> 143       columns_by_name=columns_by_name)
    144 
    145   # If the name already exists, re-use the column from columns_by_name,

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/feature_column/feature_column_v2.py in _from_config(cls, config, custom_objects, columns_by_name)
   2871     kwargs = _standardize_and_copy_config(config)
   2872     kwargs['normalizer_fn'] = generic_utils.deserialize_keras_object(
-> 2873         config['normalizer_fn'], custom_objects=custom_objects)
   2874     kwargs['dtype'] = dtypes.as_dtype(config['dtype'])
   2875 

~/anaconda3/envs/tf/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    206       obj = _GLOBAL_CUSTOM_OBJECTS[object_name]
    207     else:
--> 208       obj = module_objects.get(object_name)
    209       if obj is None:
    210         raise ValueError('Unknown ' + printable_module_name + ':' + object_name)

AttributeError: 'NoneType' object has no attribute 'get'

Answer 1

This error can be caused by having custom logic in your model, and not providing the custom logic in the custom_objects arguments when calling load_model , model_from_json , etc.

In my case, this was a function passed to layers.Lambda() :

@tf.function()
def shift_cols(x):
    """
    Transforms (32, 32, 3) to (3, 32, 32) 
    """
    return tf.transpose(x, [0, 3, 1, 2])

input_layer = Input(shape=(32, 32, 3), dtype='float32', name='features')
x = Lambda(shift_cols, input_shape=(32, 32, 3))(input_layer)  # shape=(3, 32, 32)
x = Flatten()(x) 
x = Dense(64, activation='softmax')(x) 
model = Model(inputs=input_layer, outputs={'output': x})
model.save('./model_path/')

To load the model without the specified error, you must provide the shift_cols function in the custom_objects argument in the load call:

model = tf.keras.models.load_model('./model_path/', custom_objects={
    'shift_cols': shift_cols
})

Answer 2

This is how i solve the problem . I looked at the following url:- https://github.com/tensorflow/tensorflow/issues/31927

I put my feature columns , model all in one cell and did exactly as a person did in last post in the above url. But in the above url dtype in indicator column is float but i replaced it with int32 because i was getting other error and it worked

#**************************** our main model  *********************************

keras.backend.clear_session()





feature_columns = []
feature_layer_inputs = {} 
feature_columns.append(tf.feature_column.numeric_column('col_name1'))  
feature_layer_inputs['col_name1'] = tf.keras.Input(shape=(1,), name='col_name1')  

feature_columns.append(tf.feature_column.numeric_column('col_name2')) 
feature_layer_inputs['col_name2'] = tf.keras.Input(shape=(1,), name='col_name2')  

feature_columns.append(tf.feature_column.numeric_column('col_name3'))    
feature_layer_inputs['col_name3'] = tf.keras.Input(shape=(1,), name='col_name3')  

feature_columns.append(tf.feature_column.numeric_column('col_name4'))    
feature_layer_inputs['col_name4'] = tf.keras.Input(shape=(1,), name='col_name4')  

feature_columns.append(tf.feature_column.numeric_column('col_name5'))    
feature_layer_inputs['col_name5'] = tf.keras.Input(shape=(1,), name='col_name5')  

feature_columns.append(tf.feature_column.numeric_column('col_name6'))    
feature_layer_inputs['col_name6'] = tf.keras.Input(shape=(1,), name='col_name6')  

col_name7  = tf.feature_column.categorical_column_with_identity('col_name7',10,default_value=None)
feature_columns.append(tf.feature_column.indicator_column(col_name7))
feature_layer_inputs['col_name7'] = tf.keras.Input(shape=(1,), name='col_name7', dtype= tf.int32)


col_name8  = tf.feature_column.categorical_column_with_identity('col_name8',50,default_value=None)
feature_columns.append(tf.feature_column.indicator_column(col_name8))
feature_layer_inputs['col_name8'] = tf.keras.Input(shape=(1,), name='col_name8', dtype= tf.int32)

col_name9  = tf.feature_column.categorical_column_with_identity('col_name9',12,default_value=None)
feature_columns.append(tf.feature_column.indicator_column(col_name9))
feature_layer_inputs['col_name9'] = tf.keras.Input(shape=(1,), name='col_name9', dtype= tf.int32)

feature_layer = tf.keras.layers.DenseFeatures(feature_columns)
feature_layer_outputs = feature_layer(feature_layer_inputs)


x = keras.layers.Dense(32, activation='relu')(feature_layer_outputs)
x = keras.layers.Dense(32, activation='relu')(x)
output = keras.layers.Dense(1)(x)

model = keras.Model(inputs=[v for v in feature_layer_inputs.values()], outputs=output)

model.compile(optimizer='adam',loss='mae')

model.fit(train_ds , epochs=50,verbose=1)

Now there is not error while saving or loading model