hub使用ernie tiny 训练的分类模型,为何输出的label会变化?
收藏
class TextClassification:
def __init__(self):
self.label_map = {0: '政治', 1: '军事', 2: '经济'}
self.model = hub.Module(
name='ernie_tiny',
version='2.0.1',
task='seq-cls',
load_checkpoint='./model/classification.pdparams',
label_map=self.label_map)
def __call__(self, data, max_seq_len=150, allow_empty=True):
"""
分类
:param data: 传入list数据
:param max_seq_len: 最大长度
:param allow_empty: 返回结果是否允许为空
:return:
"""
return self.model.predict(data, max_seq_len=150, batch_size=1, use_gpu=False, return_prob=True)
text_cls = TextClassification()
data = [["土耳其军工除了组装大量的F-16战斗机外,还研制了大量F-16战斗机的机载武器,近年,土耳其军工公开了自己的先进中距空空导弹和近距空空导弹,可以看到图片中两型导弹的外形,与美国军工的AIM-120C和AIM-9X空空导弹非常相似。"]]
label, probs = text_cls(data)
print(label, probs)
输出的结果:
第一次运行:
['经济'] [[0.3787701725959778, 0.2302614450454712, 0.39096832275390625]]
第二次运行:
['经济'] [[0.29642435908317566, 0.31800222396850586, 0.3855734169483185]]
模型是固定的,输入也是固定的,为什么输出的probs会变化?而且有时候会预测为其他标签。
0
收藏
请登录后评论
再一次输出,这标签就成了政治
['政治'] [[0.3992163836956024, 0.28848692774772644, 0.312296599149704]]
我看了一下,ernie tiny这个预训练模型 predict的参数中没有return_prob啊,你这个是怎么弄的?
我查看了一下新版的hub代码,里面有这个参数,然后我就从源码编译安装了一下。
你试试增大下数据量,让模型去做下微调训练后再去做预测应该能让模型稳定及效果提升些。
好的,我试试