Siamesenetworkの学習済みモデルを表示したい

Question

Kerasが公開しているSiamese Networkのモデルを保存し、表示したいと考えています。
保存まではできたのですが表示する際にエラーで表示できません。
表示させる方法を教えていただきたいです。
このLINKのものを下記のように変更しモデルを保存できるようしました。

Siamese Network

import random
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
from time import time
import os

t_start = time() # 開始時間

"""
## Hyperparameters
"""
epochs = 10
batch_size = 128
margin = 1  # Margin for constrastive loss.

(x_train_val, y_train_val), (x_test, y_test) = keras.datasets.mnist.load_data()

# Change the data type to a floating point format
x_train_val = x_train_val.astype("float32")
x_test = x_test.astype("float32")

# Keep 50% of train_val  in validation set
x_train, x_val = x_train_val[:30000], x_train_val[30000:]
y_train, y_val = y_train_val[:30000], y_train_val[30000:]
del x_train_val, y_train_val

def make_pairs(x, y):
    num_classes = max(y) + 1
    digit_indices = [np.where(y == i)[0] for i in range(num_classes)]

    pairs = []
    labels = []

    for idx1 in range(len(x)):
        # add a matching example
        x1 = x[idx1]
        label1 = y[idx1]
        idx2 = random.choice(digit_indices[label1])
        x2 = x[idx2]

        pairs += [[x1, x2]]
        labels += [1]

        # add a non-matching example
        label2 = random.randint(0, num_classes - 1)
        while label2 == label1:
            label2 = random.randint(0, num_classes - 1)

        idx2 = random.choice(digit_indices[label2])
        x2 = x[idx2]

        pairs += [[x1, x2]]
        labels += [0]

    return np.array(pairs), np.array(labels).astype("float32")


# make train pairs
pairs_train, labels_train = make_pairs(x_train, y_train)

# make validation pairs
pairs_val, labels_val = make_pairs(x_val, y_val)

# make test pairs
pairs_test, labels_test = make_pairs(x_test, y_test)

x_train_1 = pairs_train[:, 0]  # x_train_1.shape is (60000, 28, 28)
x_train_1=x_train_1.reshape(x_train_1.shape[0],x_train_1.shape[1],x_train_1.shape[2],1)
x_train_2 = pairs_train[:, 1]
x_train_2=x_train_2.reshape(x_train_2.shape[0],x_train_2.shape[1],x_train_2.shape[2],1)

"""
Split the validation pairs
"""

x_val_1 = pairs_val[:, 0]  # x_val_1.shape = (60000, 28, 28)
x_val_1=x_val_1.reshape(x_val_1.shape[0],x_val_1.shape[1],x_val_1.shape[2],1)
x_val_2 = pairs_val[:, 1]
x_val_2=x_val_2.reshape(x_val_2.shape[0],x_val_2.shape[1],x_val_2.shape[2],1)

"""
Split the test pairs
"""

x_test_1 = pairs_test[:, 0]  # x_test_1.shape = (20000, 28, 28)
x_test_1=x_test_1.reshape(x_test_1.shape[0],x_test_1.shape[1],x_test_1.shape[2],1)
x_test_2 = pairs_test[:, 1]
x_test_2=x_test_2.reshape(x_test_2.shape[0],x_test_2.shape[1],x_test_2.shape[2],1)

# Provided two tensors t1 and t2
# Euclidean distance = sqrt(sum(square(t1-t2)))
def euclidean_distance(vects):

    print(vects)
    x, y = vects
    print(x)
    print(y)
    sum_square = tf.math.reduce_sum(tf.math.square(x - y), axis=1, keepdims=True)
    return tf.math.sqrt(tf.math.maximum(sum_square, tf.keras.backend.epsilon()))


input = layers.Input((28,28,1))
print(input.shape)
x = tf.keras.layers.BatchNormalization()(input)
x = layers.Conv2D(4, (5, 5), activation="tanh")(x)
x = layers.AveragePooling2D(pool_size=(2, 2))(x)
x = layers.Conv2D(16, (5, 5), activation="tanh")(x)
x = layers.AveragePooling2D(pool_size=(2, 2))(x)
x = layers.Flatten()(x)

x = tf.keras.layers.BatchNormalization()(x)
x = layers.Dense(10, activation="tanh")(x)
print(x.shape)
embedding_network = keras.Model(input, x)


input_1 = layers.Input((28,28,1))
print(input_1.shape)
input_2 = layers.Input((28,28,1))
print(input_2.shape)

tower_1 = embedding_network(input_1)
tower_2 = embedding_network(input_2)

merge_layer = layers.Lambda(euclidean_distance)([tower_1, tower_2])
normal_layer = tf.keras.layers.BatchNormalization()(merge_layer)
output_layer = layers.Dense(1, activation="sigmoid")(normal_layer)
siamese = keras.Model(inputs=[input_1, input_2], outputs=output_layer)


"""
## Define the constrastive Loss
"""


def loss(margin=1):

    # Contrastive loss = mean( (1-true_value) * square(prediction) +
    #                         true_value * square( max(margin-prediction, 0) ))
    def contrastive_loss(y_true, y_pred):

        square_pred = tf.math.square(y_pred)
        margin_square = tf.math.square(tf.math.maximum(margin - (y_pred), 0))
        return tf.math.reduce_mean(
            (1 - y_true) * square_pred + (y_true) * margin_square
        )

    return contrastive_loss


"""
## Compile the model with the contrastive loss
"""

siamese.compile(loss=loss(margin=margin), optimizer="RMSprop", metrics=["accuracy"])
tf.keras.utils.plot_model(siamese, to_file = 'model2.png', show_shapes = True, show_layer_names = True)

siamese.summary()



"""
## Train the model
"""

history = siamese.fit(
    [x_train_1, x_train_2],
    labels_train,
    validation_data=([x_val_1, x_val_2], labels_val),
    batch_size=batch_size,
    epochs=epochs
)

siamese.save('/home/practice/siamese_model/model.hdf5')




t_end = time() #終了時間
t_elapsed = t_end - t_start

print("処理時間は{0}".format(t_elapsed))

学習済みモデルを表示させるプログラム

#from sklearn.externals import joblib
import argparse, os, glob, datetime, cv2, re, sys, math
import tensorflow as tf
from tensorflow.keras import backend as K
import random as rand
import numpy as np
import keras
from keras.layers import Lambda
from keras.layers import  Input, Conv2D, BatchNormalization, Activation, Subtract
from tensorflow.keras.models import Model, load_model
from keras.losses import categorical_crossentropy
from keras.optimizers import Adam
from PIL import Image
from pathlib import Path
from time import time

def contrastive_loss(y_true, y_pred):
    margin = 1
    square_pred = K.square(y_pred)
    margin_square = K.square(K.maximum(margin - y_pred, 0))
    return K.mean(y_true * square_pred + (1 - y_true) * margin_square)

dir = '/home/practice/siamese_model/model.hdf5'


base_model = load_model(dir,custom_objects={'contrastive_loss': contrastive_loss,'Lambda':Lambda},compile=False)

base_model.summary()

発生しているエラー

Traceback (most recent call last):
  File "model_loader.py", line 28, in <module>
    base_model = load_model(dir,custom_objects={'contrastive_loss': contrastive_loss,'Lambda':Lambda},compile=False)
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/save.py", line 146, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/hdf5_format.py", line 212, in load_model_from_hdf5
    custom_objects=custom_objects)
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/model_config.py", line 55, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/layers/serialization.py", line 89, in deserialize
    printable_module_name='layer')
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/utils/generic_utils.py", line 192, in deserialize_keras_object
    list(custom_objects.items())))
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1121, in from_config
    process_layer(layer_data)
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1105, in process_layer
    layer = deserialize_layer(layer_data, custom_objects=custom_objects)
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/layers/serialization.py", line 89, in deserialize
    printable_module_name='layer')
  File "/home/.local/lib/python3.6/site-packages/tensorflow/python/keras/utils/generic_utils.py", line 192, in deserialize_keras_object
    list(custom_objects.items())))
  File "/home/.local/lib/python3.6/site-packages/keras/layers/core.py", line 793, in from_config
    return cls(**config)
  File "/home/.local/lib/python3.6/site-packages/keras/legacy/interfaces.py", line 91, in wrapper
    return func(*args, **kwargs)
  File "/home/.local/lib/python3.6/site-packages/keras/layers/core.py", line 649, in __init__
    super(Lambda, self).__init__(**kwargs)
  File "/home/.local/lib/python3.6/site-packages/keras/engine/base_layer.py", line 128, in __init__
    raise TypeError('Keyword argument not understood:', kwarg)
TypeError: ('Keyword argument not understood:', 'module')