本文主要对加载模型报错信息进行介绍。
一、错误描述
通过 torch.load 函数加载模型文件时出错了,报错信息如下。
Traceback (most recent call last):
File "cli_demo.py", line 10, in <module>
model = AutoModel.from_pretrained("D:\code\ChatGLM-6B-main\THUDM\chatglm2-6b", trust_remote_code=True).float()
File "D:\Software\anaconda3\envs\chat6b\lib\site-packages\transformers\models\auto\auto_factory.py", line 550, in from_pretrained
model_class = get_class_from_dynamic_module(
File "D:\Software\anaconda3\envs\chat6b\lib\site-packages\transformers\dynamic_module_utils.py", line 510, in get_class_from_dynamic_module
return get_class_in_module(class_name, final_module)
File "D:\Software\anaconda3\envs\chat6b\lib\site-packages\transformers\dynamic_module_utils.py", line 208, in get_class_in_module
module_spec.loader.exec_module(module)
File "<frozen importlib._bootstrap_external>", line 843, in exec_module
File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
File "C:\Users\86188\.cache\huggingface\modules\transformers_modules\chatglm2-6b\modeling_chatglm.py", line 958, in <module>
securityCheck = SecurityChecking()
File "C:\Users\86188\.cache\huggingface\modules\transformers_modules\chatglm2-6b\modeling_chatglm.py", line 933, in __init__
self.securityCheckModel = SecurityCheckModel()
File "C:\Users\86188\.cache\huggingface\modules\transformers_modules\chatglm2-6b\modeling_chatglm.py", line 861, in __init__
self.model = torch.load("D:\code\ChatGLM-6B-main\THUDM\chatglm2-6b\model\BiLSTM0723.pt")
File "D:\Software\anaconda3\envs\chat6b\lib\site-packages\torch\serialization.py", line 1014, in load
return _load(opened_zipfile,
File "D:\Software\anaconda3\envs\chat6b\lib\site-packages\torch\serialization.py", line 1422, in _load
result = unpickler.load()
File "D:\Software\anaconda3\envs\chat6b\lib\site-packages\torch\serialization.py", line 1415, in find_class
return super().find_class(mod_name, name)
AttributeError: Can't get attribute 'BiLSTMClassifier' on <module '__main__' from 'cli_demo.py'>
大致意思是说在 cli_demo.py 中没有找到 BiLSTMClassifier 类,BiLSTMClassifier 类即为 torch.load 加载模型的类,下面具体分析下原因和解决方法。
二、原因分析
因为在加载模型时,当前运行的脚本无法找到模型中使用的类 (本文是BiLSTMClassifier
) 的定义。当使用 torch.save 保存一个模型时,模型的类名会被序列化。在加载模型时,PyTorch 需要通过类名找到对应的类定义。如果当前脚本中没有导入或定义该类,就会出现这个错误。
三、解决方法
解决方法就是拷贝类的定义到当前文件中(本文是 cli_demo.py,见报错信息)或者 导入模型类到当前文件中,两种可以分别尝试下。
导入类可以通过如下方式:
from BiLSTM import BiLSTMClassifier
# from 文件 import 类名