如何加载ktrain训练的预训练BERT模型？加载报错求助

用ktrain训练BERT文本分类模型并完成保存，训练时运行正常，保存predictor后生成包含tf_model.h5和tf_model.preproc的bert_model文件夹，但加载该predictor时出现报错。

训练及保存代码 #

y_train = [encoding[x] for x in y_train]
y_test = [encoding[x] for x in y_test]
(x_train,  y_train), (x_test, y_test), preproc = text.texts_from_array(x_train=X_train, y_train=y_train,
                                                                       x_test=X_test, y_test=y_test,
                                                                       class_names=class_names,
                                                                       preprocess_mode='bert',
                                                                       maxlen=350,
                                                                       max_features=35000)

model = text.text_classifier(
    'bert', train_data=(x_train, y_train), preproc=preproc)
learner = ktrain.get_learner(model, train_data=(x_train, y_train),
                             val_data=(x_test, y_test),
                             batch_size=6)
learner.fit_onecycle(2e-5, 3)
learner.validate(val_data=(x_test, y_test), class_names=class_names)
predictor = ktrain.get_predictor(learner.model, preproc)
predictor.get_classes()

import time
message = 'she left me while i was waiting for her and she isnt herself '

start_time = time.time()
prediction = predictor.predict(message)

print('predicted: {} ({:.2f})'.format(prediction, (time.time() - start_time)))
predictor.save("models/bert_model")

加载代码 #

from keras.models import load_model
import ktrain
from ktrain import text
import time

message = 'she left me while i was waiting for her and she isnt herself '
model_path = "/home/priyanshi/Desktop/folder/models/bert_model/"

predictor = ktrain.load_predictor(model_path)

start_time = time.time()
prediction = predictor.predict(message)

print('predicted: {} ({:.2f} seconds)'.format(
    prediction, (time.time() - start_time)))

报错信息 #

python -u "/home/priyanshi/Desktop/folder/test2.py"
2024-06-02 00:59:18.984734: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-06-02 00:59:18.987632: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-06-02 00:59:19.031606: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-06-02 00:59:19.734612: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Call to keras.models.load_model failed. Try manually invoking this function to investigate error and report issue if necessary.
Traceback (most recent call last):
File "/home/priyanshi/.local/lib/python3.10/site-packages/ktrain/core.py", line 2048, in _load_model
model = keras.models.load_model(
File "/home/priyanshi/.local/lib/python3.10/site-packages/tf_keras/src/saving/saving_api.py", line 262, in load_model
return legacy_sm_saving_lib.load_model(
File "/home/priyanshi/.local/lib/python3.10/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/home/priyanshi/.local/lib/python3.10/site-packages/tensorflow/python/saved_model/loader_impl.py", line 119, in parse_saved_model
raise IOError(
OSError: SavedModel file does not exist at: /home/priyanshi/Desktop/folder/models/bert_model//{saved_model.pbtxt|saved_model.pb}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/priyanshi/Desktop/folder/test2.py", line 13, in 
predictor = ktrain.load_predictor(model_path)
File "/home/priyanshi/.local/lib/python3.10/site-packages/ktrain/core.py", line 1848, in load_predictor
model = _load_model(fpath, preproc=preproc, custom_objects=custom_objects)
File "/home/priyanshi/.local/lib/python3.10/site-packages/ktrain/core.py", line 2055, in _load_model
raise Exception("Error detected: %s" % (e))
Exception: Error detected: SavedModel file does not exist at: /home/priyanshi/Desktop/folder/models/bert_model//{saved_model.pbtxt|saved_model.pb}

解决方法 #

• 指定模型文件名：ktrain.load_predictor默认会寻找SavedModel格式文件，而你保存的是H5格式，需要手动指定模型文件名：

  predictor = ktrain.load_predictor(model_path, model_filename='tf_model.h5')
  

• 统一环境版本：确保训练和加载环境的ktrain、TensorFlow版本完全一致，版本不兼容可能导致模型格式识别错误。
• 检查路径正确性：移除model_path末尾的斜杠，避免路径拼接出现重复斜杠问题：

  model_path = "/home/priyanshi/Desktop/folder/models/bert_model"
  

• 验证模型文件完整性：手动加载H5模型排查文件是否损坏：

  from tensorflow.keras.models import load_model
  model = load_model('/home/priyanshi/Desktop/folder/models/bert_model/tf_model.h5')
  

内容的提问来源于stack exchange，提问作者Priyanshi Soni