使用pickle加载Colab训练的NLP模型失败求助

问题场景 #

在Google Colab训练NLP模型（50个epochs，耗时约3小时），使用pickle.dump将模型保存为model.pkl并下载到本地后，执行pickle.load加载时出现如下报错：

ValueError                                Traceback (most recent call last)
~\AppData\Local\Temp\ipykernel_15292\2076388884.py in ()
----> 1 model = pickle.load(open('model.pkl','rb'))

c:\Users\balaj\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\saving\pickle_utils.py in deserialize_model_from_bytecode(serialized_model)
     46         model = saving_lib.load_model(filepath, safe_mode=False)
     47     except Exception as e:
---> 48         raise e
     49     else:
     50         return model

c:\Users\balaj\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\saving\pickle_utils.py in deserialize_model_from_bytecode(serialized_model)
     44         # serialized as a string by Dense.get_config()) will require
     45         # a custom_object_scope.
---> 46         model = saving_lib.load_model(filepath, safe_mode=False)
     47     except Exception as e:
     48         raise e

c:\Users\balaj\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\saving\saving_lib.py in load_model(filepath, custom_objects, compile, safe_mode)
    275 
    276     except Exception as e:
---> 277         raise e
    278     else:
    279         return model
...
---> 256                 raise ValueError(
    257                     f"Expected a {_VARS_FNAME}.h5 or {_VARS_FNAME}.npz file."
    258                 )

ValueError: Expected a model.weights.h5 or model.weights.npz file.

问题原因 #

Keras模型的结构与权重是分离存储的，pickle仅能序列化模型的字节码，无法完整保存Keras模型的权重关联信息，导致本地加载时无法找到对应的权重文件，从而抛出错误。

无需重新训练的解决办法 #

方案1：仍可访问Colab训练环境（模型还在内存中） #

直接用Keras官方推荐的方法重新保存模型，替换原pickle保存方式：

1. 在Colab中执行以下代码保存模型：

   # 保存为HDF5格式（单文件）
   model.save("nlp_model.h5")
   
   # 或保存为SavedModel格式（推荐，跨版本兼容性更好）
   model.save("nlp_saved_model")
   

2. 将生成的nlp_model.h5文件（或nlp_saved_model文件夹）下载到本地
3. 本地加载模型：

   from keras.models import load_model
   
   # 加载HDF5模型
   model = load_model("nlp_model.h5")
   
   # 或加载SavedModel格式
   model = load_model("nlp_saved_model")
   

方案2：无法访问Colab环境，仅持有model.pkl文件 #

尝试通过Keras内置工具修复加载，或手动重建模型结构后加载权重：

1. 尝试用Keras的pickle_utils反序列化：

   import pickle
   from keras.saving.pickle_utils import deserialize_model_from_bytecode
   
   with open('model.pkl', 'rb') as f:
       serialized_data = pickle.load(f)
   try:
       model = deserialize_model_from_bytecode(serialized_data)
   except Exception as e:
       print(f"直接反序列化失败: {e}")
       # 进入手动重建流程
   

2. 手动重建模型结构并加载权重：

   注意：必须完全复刻训练时的模型结构（包括层类型、参数、顺序、激活函数等），否则权重加载会失败

   from keras.models import Sequential
   from keras.layers import Embedding, LSTM, Dense
   import h5py
   import io
   
   # 示例：重建和训练时一致的模型结构（请替换为你的实际结构）
   vocab_size = 10000  # 替换为你的实际词汇表大小
   model = Sequential()
   model.add(Embedding(input_dim=vocab_size, output_dim=128))
   model.add(LSTM(64))
   model.add(Dense(1, activation='sigmoid'))
   
   # 从pickle文件中提取权重并加载
   with open('model.pkl', 'rb') as f:
       serialized_data = pickle.load(f)
   if isinstance(serialized_data, bytes):
       with h5py.File(io.BytesIO(serialized_data), 'r') as weight_file:
           model.load_weights(weight_file)
   

后续注意事项 #

• 保存Keras模型时，优先使用model.save()方法，避免直接用pickle，这是Keras官方明确不推荐的做法，容易引发跨环境兼容性问题。
• 若需跨设备/版本迁移模型，优先选择SavedModel格式，其兼容性优于HDF5格式。

内容的提问来源于stack exchange，提问作者BALAJI P RA2111026010310