bert4keras模型转onnx

bert4keras模型转onnx

https://kexue.fm/archives/8454

#! -*- coding: utf-8 -*-
import os
os.environ['TF_KERAS'] = '1'
from bert4keras.backend import keras
from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer

'''
生成saved-model的时候
bert4keras==0.10.6
keras==2.3.1
tensorflow-gpu==1.15.4
tensorflow-hub==0.12.0
h5py==2.10.0
'''
'''
生成后将saved-model转onnx，我是直接pip install tensorflow==2.5.0
建议创建一个新环境
bert4keras==0.10.6
keras==2.3.1
tensorflow==2.5.0
h5py==3.1.0
keras-nightly==2.5.0.dev2021032900
tensorflow-estimator==2.5.0
tf2onnx==1.9.1
onnx==1.9.0
onnxruntime==1.8.0
然后
# python -m tf2onnx.convert --saved-model encoder_model_tf --output encoder_simbert.onnx --opset 13
# python -m tf2onnx.convert --saved-model generate_model_tf --output generate_simbert.onnx --opset 13
'''

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# 基本信息
maxlen = 128

# 模型配置
model_file = 'chinese_roformer-sim-char_L-12_H-768_A-12'
config_path = '{}/bert_config.json'.format(model_file)
checkpoint_path = '{}/bert_model.ckpt'.format(model_file)
# checkpoint_path = None
dict_path = '{}/vocab.txt'.format(model_file)

# 建立分词器
tokenizer = Tokenizer(dict_path, do_lower_case=True)

roformer = build_transformer_model(
    config_path,
    checkpoint_path,
    model='roformer',
    application='unilm',
    with_pool='linear'
)
# 可以加载自己训练后的模型
output_best_model = 'model_path/best_model.weights'
if checkpoint_path is None:
    roformer.load_weights(output_best_model)

# 向量生成模型
# encoder = keras.models.Model(roformer.inputs, roformer.outputs[0])
# encoder.save('model_path/encoder_model_tf', save_format='tf')

# 解码器模型
seq2seq = keras.models.Model(roformer.inputs, roformer.outputs[1])
outputs = [
    keras.layers.Lambda(lambda x: x[:, -1])(output)
    for output in seq2seq.outputs
]
generate_model = keras.models.Model(seq2seq.inputs, outputs)
generate_model.save('model_path/generate_model_tf', save_format='tf')

# import keras2onnx
# onnx_model = keras2onnx.convert_keras(encoder)
# keras2onnx.save_model(onnx_model, 'bert-sim.onnx')

得到 saved_model 后再转onnx

'''
经过验证这个时候需要tensorflow==2.5.0  
安装pip install tensorflow==2.5.0  
'''
python -m tf2onnx.convert --saved-model encoder_model_tf --output simbert.onnx --opset 13

# 预测的时候还是用的tensorflow-gpu==1.15.4
如果是在一个环境里面操作，这里装两个tensorflow后再用bert4keras可能会出现问题，
uninstall tensorflow 后再install; h5py的版本也要注意。
pip uninstall keras-nightly
pip uninstall tensorflow
pip uninstall -y tensorflow-gpu
pip install keras==2.3.1
pip install tensorflow-gpu==1.15.4
pip install h5py==2.10.0
--opset 10 也是能保存成功，但是预测的时候出错了。

用Netron查看simbert.onnx模型，

INPUTS
name: Input-Segment
type: float32[unk__3503,unk__3504]

name: Input-Token
type: float32[unk__3505,unk__3506]

import onnxruntime
import numpy as np
from model_path import tokenization
from model_path.bert_input_process import *

class SimBertONNX:
    def __init__(self, weights, max_seg_len=128, vocab_file='chinese_roformer-sim-char_L-12_H-768_A-12/vocab.txt'):
        self.sess = onnxruntime.InferenceSession(weights)
        self.max_seg_len = max_seg_len
        self.tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=True)

    def predict(self, sent):
        feature = format_ner_input(sent, max_seq_length=self.max_seg_len, tokenizer=self.tokenizer)
        token = [i for i in feature['input_ids'][0] if i != 0]
        input_dict = {'Input-Token': [token],  "Input-Segment": [[0]*len(token)]}
        # {'Input-Token': [[101, 4636, 1668, 4686, 6994, 102]], 'Input-Segment': [[0, 0, 0, 0, 0, 0]]}
        result = self.sess.run(None, input_dict)[0][0]
        return result


pb_path = "model_path/simbert13.onnx"
simbert = SimBertONNX(pb_path)
print(simbert.predict('科学空间'))
[-1.11289009e-01 -5.35200179e-01  1.08341195e-01 -1.98119685e-01
 -5.59085608e-02  4.03575003e-01  1.67680249e-01 -2.41547719e-01
  4.21358049e-01 -2.02915221e-01 -1.85670257e-01 -2.22553447e-01
...

# bert_input_process

import numpy as np

def _truncate_seq_pair(tokens_a, tokens_b, max_length):
    while True:
        total_length = len(tokens_a) + len(tokens_b)
        if total_length <= max_length:
            break
        if len(tokens_a) > len(tokens_b):
            tokens_a.pop()
        else:
            tokens_b.pop()


class InputFeatures(object):
    def __init__(self, input_ids, input_mask, segment_ids, label_id, is_real_example=True):
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.label_id = label_id
        self.is_real_example = is_real_example


class PaddingInputExample(object):
    pass


class InputExample(object):
    def __init__(self, guid, text_a, text_b=None, label=None):
        self.guid = guid
        self.text_a = text_a
        self.text_b = text_b
        self.label = label


def convert_ner_example(example, max_seq_length, tokenizer):
    if isinstance(example, PaddingInputExample):
        return InputFeatures(
            input_ids=[0] * max_seq_length,
            input_mask=[0] * max_seq_length,
            segment_ids=[0] * max_seq_length,
            label_id=[0] * max_seq_length,
            is_real_example=False)

    tokens_a = tokenizer.tokenize(example.text_a)
    tokens_b = None
    if example.text_b:
        tokens_b = tokenizer.tokenize(example.text_b)

    if tokens_b:
        _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)
    else:
        if len(tokens_a) > max_seq_length - 2:
            tokens_a = tokens_a[0:(max_seq_length - 2)]
    tokens = []
    segment_ids = []
    label_id = [0] * max_seq_length
    tokens.append("[CLS]")
    segment_ids.append(0)
    for token in tokens_a:
        tokens.append(token)
        segment_ids.append(0)
    tokens.append("[SEP]")
    segment_ids.append(0)

    if tokens_b:
        for token in tokens_b:
            tokens.append(token)
            segment_ids.append(1)
        tokens.append("[SEP]")
        segment_ids.append(1)

    input_ids = tokenizer.convert_tokens_to_ids(tokens)
    label_id[0] = input_ids[0]
    label_id[len(tokens) - 1] = input_ids[len(tokens) - 1]

    input_mask = [1] * len(input_ids)

    while len(input_ids) < max_seq_length:
        input_ids.append(0)
        input_mask.append(0)
        segment_ids.append(0)

    assert len(input_ids) == max_seq_length
    assert len(input_mask) == max_seq_length
    assert len(segment_ids) == max_seq_length

    feature = InputFeatures(
        input_ids=input_ids,
        input_mask=input_mask,
        segment_ids=segment_ids,
        label_id=label_id,
        is_real_example=True)
    return feature


def format_ner_input(text_a, text_b=None, gid="id", max_seq_length=128, tokenizer=None):
    predict_example = InputExample(gid, text_a, text_b, None)
    feature = convert_ner_example(predict_example, max_seq_length, tokenizer)
    return {
        "input_ids": [feature.input_ids],
        "input_mask": [feature.input_mask],
        "segment_ids": [feature.segment_ids],
        "label_ids": [feature.label_id],
    }


def format_simbert_input(text_a_list, max_seq_length=128, tokenizer=None):
    input_ids = []
    segment_ids = []
    for text_a in text_a_list:
        predict_example = InputExample(guid="id", text_a=text_a, text_b=None, label=None)
        feature = convert_ner_example(predict_example, max_seq_length, tokenizer)
        input_ids.append(feature.input_ids)
        segment_ids.append(feature.segment_ids)
    return input_ids, segment_ids


def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    x = np.exp(x)
    return x / x.sum(axis=axis, keepdims=True)


class AutoRegressiveDecoder(object):
    def __init__(self, start_id, end_id, maxlen, minlen=1):
        self.start_id = start_id
        self.end_id = end_id
        self.maxlen = maxlen
        self.minlen = minlen
        self.models = {}
        if start_id is None:
            self.first_output_ids = np.empty((1, 0), dtype=int)
        else:
            self.first_output_ids = np.array([[self.start_id]])

    @staticmethod
    def wraps(default_rtype='probas', use_states=False):
        def actual_decorator(predict):
            def new_predict(self, inputs, output_ids, states, temperature=1, rtype=default_rtype):
                assert rtype in ['probas', 'logits']
                prediction = predict(self, inputs, output_ids, states)
                if not use_states:
                    prediction = (prediction, None)
                if default_rtype == 'logits':
                    prediction = (softmax(prediction[0] / temperature), prediction[1])
                elif temperature != 1:
                    probas = np.power(prediction[0], 1.0 / temperature)
                    probas = probas / probas.sum(axis=-1, keepdims=True)
                    prediction = (probas, prediction[1])
                if rtype == 'probas':
                    return prediction
                else:
                    return np.log(prediction[0] + 1e-12), prediction[1]

            return new_predict

        return actual_decorator

    def predict(self, inputs, output_ids, states=None):
        raise NotImplementedError

    def random_sample(self, inputs, n, topp=None, states=None, temperature=1, min_ends=1):
        inputs = [np.array([i]) for i in inputs]
        output_ids = self.first_output_ids
        results = []
        for step in range(self.maxlen):
            probas, states = self.predict(
                inputs, output_ids, states, temperature, 'probas'
            )  # 计算当前概率
            probas /= probas.sum(axis=1, keepdims=True)  # 确保归一化
            if step == 0:  # 第1步预测后将结果重复n次
                probas = np.repeat(probas, n, axis=0)
                inputs = [np.repeat(i, n, axis=0) for i in inputs]
                output_ids = np.repeat(output_ids, n, axis=0)
            if topp is not None:
                p_indices = probas.argsort(axis=1)[:, ::-1]  # 从高到低排序
                probas = np.take_along_axis(probas, p_indices, axis=1)  # 排序概率
                cumsum_probas = np.cumsum(probas, axis=1)  # 累积概率
                flag = np.roll(cumsum_probas >= topp, 1, axis=1)  # 标记超过topp的部分
                flag[:, 0] = False  # 结合上面的np.roll，实现平移一位的效果
                probas[flag] = 0  # 后面的全部置零
                probas /= probas.sum(axis=1, keepdims=True)  # 重新归一化
            sample_func = lambda p: np.random.choice(len(p), p=p)  # 按概率采样函数
            sample_ids = np.apply_along_axis(sample_func, 1, probas)  # 执行采样
            sample_ids = sample_ids.reshape((-1, 1))  # 对齐形状
            if topp is not None:
                sample_ids = np.take_along_axis(p_indices, sample_ids, axis=1)
            output_ids = np.concatenate([output_ids, sample_ids], 1)  # 更新输出
            end_counts = (output_ids == self.end_id).sum(1)  # 统计出现的end标记
            if output_ids.shape[1] >= self.minlen:  # 最短长度判断
                flag = (end_counts == min_ends)  # 标记已完成序列
                if flag.any():  # 如果有已完成的
                    for ids in output_ids[flag]:  # 存好已完成序列
                        results.append(ids)
                    flag = (flag == False)  # 标记未完成序列
                    inputs = [i[flag] for i in inputs]  # 只保留未完成部分输入
                    output_ids = output_ids[flag]  # 只保留未完成部分候选集
                    end_counts = end_counts[flag]  # 只保留未完成部分end计数
                    if len(output_ids) == 0:
                        break
        # 如果还有未完成序列，直接放入结果
        for ids in output_ids:
            results.append(ids)
        # 返回结果
        return results

感谢科学空间的分享