tensorflow-gpu 2.0.0中load_model()方法的报错和解决

调用Kears中load_model()方法遇到的问题和解决方法

运行环境：tensorflow-gpu 2.0.0a0，显卡配置：GTX1050Ti
报错：Unknown entry in loss dictionary: class_name. Only expected following keys:[‘dense_2’]

程序流程如下

执行以下代码进行网络训练，调用keras自带fit函数

model = build_model()

# Train
history = model.fit(img, label, batch_size=16,
                    epochs=5, validation_split=0.05)

其中，使用的部分网络参数如下

# configure the model
model.compile(optimizer=keras.optimizers.Adam(),
              loss=keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

使用gpu训练，训练完成的模型可以在CPU版本的tensorflow 2.0.0通过load_model()正常载入和运行。
在gpu版本中调用以下函数

def load_pretrained_model(filename):
    pretrained_model = load_model(filename)
    return pretrained_model

出现报错：Unknown entry in loss dictionary: class_name. Only expected following keys:[‘dense_2’]

解决方案

尝试过其他博客中提出过的方法：把keras.losses.SparseCategoricalCrossentropy()改为keras.losses.sparse_categorical_crossentropy，依然报错。

最后通过使用另一种方式读取模型解决问题：只读取训练好的权重参数，再载入到网络中即可。

代码如下：

在保存网络时候只保存权重：

# save weights only
model.save_weights('weights.h5')
print("Saved as 'weights.h5'")

读取时先载入模型，然后载入权重：

# load weights only
pretrained_model = build_model()
pretrained_model.load_weights('weights.h5')

其中，build_model()方法如下：

# based on VGG-16
model = keras.Sequential()
# ...construct network here...
# model.add(...)

# configure the model
model.compile(optimizer=keras.optimizers.Adam(),
              loss=keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])
model.summary()
return model

最后使用

pretrained_model.predict(target)

进行预测即可。