Search code examples
pythontensorflowkerasdeep-learningkeras-layer

load_model and Lamda layer in Keras


How to load model that have lambda layer?

Here is the code to reproduce behaviour:

MEAN_LANDMARKS = np.load('data/mean_shape_68.npy')

def add_mean_landmarks(x):
    mean_landmarks = np.array(MEAN_LANDMARKS, np.float32)
    mean_landmarks = mean_landmarks.flatten()
    mean_landmarks_tf = tf.convert_to_tensor(mean_landmarks)
    x = x + mean_landmarks_tf
    return x

def get_model():
    inputs = Input(shape=(8, 128, 128, 3))
    cnn = VGG16(include_top=False, weights='imagenet', input_shape=(128, 128, 3))
    x = TimeDistributed(cnn)(inputs)
    x = TimeDistributed(Flatten())(x)
    x = LSTM(256)(x)
    x = Dense(68 * 2, activation='linear')(x)

    x = Lambda(add_mean_landmarks)(x)

    model = Model(inputs=inputs, outputs=x)
    optimizer = Adadelta()
    model.compile(optimizer=optimizer, loss='mae')

    return model

Model compiles and I can save it, but when I tried to load it with load_model function I get an error:

in add_mean_landmarks
    mean_landmarks = np.array(MEAN_LANDMARKS, np.float32)
NameError: name 'MEAN_LANDMARKS' is not defined

Аs I understand MEAN_LANDMARKS is not incorporated in graph as constant tensor. Also it's related to this question: How to add constant tensor in Keras?


Solution

  • You need to pass custom_objects argument to load_model function:

    model = load_model('model_file_name.h5', custom_objects={'MEAN_LANDMARKS': MEAN_LANDMARKS})
    

    Look for more info in Keras docs: Handling custom layers (or other custom objects) in saved models .