使用Flask部署机器学习模型时遇embedding层加载错误求解决

问题背景 #

部署基于Flask的机器学习模型时，模型加载环节触发以下错误：

ValueError: Layer 'embedding' expected 1 variables, but received 0 variables during loading. Expected: ['embedding/embeddings:0']

已尝试使用pickle、h5、joblib三种格式加载模型，均未解决问题。

现有部署代码 #

from flask import Flask, render_template, request, jsonify
import numpy as np
import re
from keras.utils.data_utils import pad_sequences
# from tensorflow.keras.models import load_model
from keras.preprocessing.text import Tokenizer
import pickle
import json

import joblib

# Load the model using joblib
model = joblib.load('model.pkl')


app = Flask(__name__)


def tokenizer_from_json(tokenizer_json):
    tokenizer = Tokenizer()
    tokenizer.word_index = tokenizer_json['word_index']
    tokenizer.index_word = {int(i): w for i, w in tokenizer_json['index_word'].items()}
    tokenizer.document_count = tokenizer_json['document_count']
    tokenizer.word_docs = tokenizer_json['word_docs']
    return tokenizer

# with open('model.pkl', 'rb') as file:
#     model = pickle.load(file)

with open('tokenizer.json', 'r') as file:
    tokenizer_json = json.load(file)
tokenizer = tokenizer_from_json(tokenizer_json)

with open('username_mapping.json', 'r') as file:
    username_mapping = json.load(file)

with open('party_mapping.json', 'r') as file:
    party_mapping = json.load(file)

max_sequence_length = 200

@app.route('/home')
def home():
    return render_template('home.html')

@app.route('/predicted', methods=['POST'])
def predicted():
    sentence = request.form['sentence']
    
    sentence = re.sub('[^a-zA-Z]', ' ', sentence)
    sentence = sentence.lower()
    tokens = sentence.split()
    tokens_sequence = tokenizer.texts_to_sequences([tokens])
    tokens_padded = pad_sequences(tokens_sequence, maxlen=max_sequence_length)
    username_index = np.argmax(model.predict(tokens_padded))
    predicted_username = list(username_mapping.keys())[list(username_mapping.values()).index(username_index)]
    predicted_party = party_mapping.get(predicted_username, 'Unknown Party')
    
    return render_template('predicted.html', party=predicted_party)

if __name__ == '__main__':
    app.run()

注：user_mapping和party_mapping为Python字典，已保存为JSON文件，在Flask代码中可正常加载使用。

修复方案 #

1. 使用Keras官方方式保存与加载模型 #

Keras模型不推荐用pickle/joblib保存，官方标准流程如下：

• 重新保存模型（训练环境执行）：

  from tensorflow.keras.models import save_model
  # 假设训练好的模型变量名为model
  save_model(model, 'my_model.h5')  # 保存为h5格式
  # 或推荐使用SavedModel格式
  save_model(model, 'my_saved_model')
  

• 在Flask中加载模型：
  替换现有模型加载代码为：

  from tensorflow.keras.models import load_model
  # 加载h5格式模型
  model = load_model('my_model.h5')
  # 或加载SavedModel格式
  # model = load_model('my_saved_model')
  

2. 检查Embedding层权重完整性 #

训练完成后，确认Embedding层权重不为空：

print(model.get_layer('embedding').get_weights())

若权重为空，说明训练过程中Embedding层未正确初始化或训练，需重新训练模型。

3. 对齐训练与部署环境的版本 #

Keras/TensorFlow版本不一致会导致模型变量结构不匹配，执行以下命令检查版本：

# 训练环境
pip show tensorflow keras
# 部署环境
pip show tensorflow keras

将部署环境的版本调整为与训练环境完全一致。

4. 修正Tokenizer加载逻辑 #

使用Keras官方提供的tokenizer_from_json方法，避免手动构造遗漏参数：

from tensorflow.keras.preprocessing.text import tokenizer_from_json

with open('tokenizer.json', 'r') as file:
    tokenizer = tokenizer_from_json(file.read())

内容的提问来源于stack exchange，提问作者Shamik Rana