fasterseq 0.0.4
fasterseq tensorflow [bert transformer] ckpt模型转换pb模型, 支持fastertransformer,加速bert等模型推理
License: Apache 2.0
Requires Python: >=3, <4
Keywords: fasterseq,bert,tensorflow bert,tensorflow transformer,seq

fasterseq tensorflow [bert transformer] ckpt模型转换pb模型, 支持fastertransformer,加速bert等模型推理

```py
# -*- coding: utf-8 -*-
import os
import sys

import tensorflow as tf
import fasterseq

'''
前言: 
        fasterseq 用于将tensorflow模型 ckpt转换pb模型
        当前转换支持普通模型转换和fastertransformer转换
        转换好的模型均可以通过nn-sdk 加速推理
        如果使用 fastertransformer 转换模型，目前只支持linux tensorflow 1.15.
        推荐 tensorflow 连接如下 ， 建议使用cuda11.3.1 环境tensorflow 1.15
        tensorflow链接: https://pan.baidu.com/s/1PXelYOJ2yqWfWfY7qAL4wA 提取码: rpxv 复制这段内容后打开百度网盘手机App，操作更方便哦
'''
ready_config = {
    "use_fastertransformer": 1, # 0 禁用fastertransormer , 1 启用fastertransormer
    "fastertransformer":{
        "cuda_version": "11.3", # 当前支持 10.2, 11.3
        "floatx": "float32",# float16, float32
        "int8_mode": 0,
        "remove_padding": False
    }
}
#初始化fastseq模块
BertModel_module = fasterseq.ready(ready_config)
if BertModel_module is None:
    raise Exception('fasterseq ready failed')

def create_model(bert_config, is_training, input_ids, input_mask, segment_ids, num_labels, use_one_hot_embeddings):
    """Creates a classification model."""
    model = BertModel_module.BertModel(
      config=bert_config,
      is_training=is_training,
      input_ids=input_ids,
      input_mask=input_mask,
      token_type_ids=segment_ids,
      use_one_hot_embeddings=use_one_hot_embeddings)

    output_layer = model.get_pooled_output()

    hidden_size = output_layer.shape[-1].value

    output_weights = tf.get_variable(
      "output_weights", [num_labels, hidden_size],
      dtype="float32",
      initializer=tf.truncated_normal_initializer(stddev=0.02))

    output_bias = tf.get_variable(
      "output_bias", [num_labels],
      dtype="float32",
      initializer=tf.zeros_initializer())

    logits = tf.matmul(output_layer, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    probabilities = tf.nn.softmax(logits, axis=-1)
    return probabilities

bert_dir=r'/data/train_data/pre_models/tf/bert/uncased_L-12_H-768_A-12'
config_file = os.path.join(bert_dir, 'bert_config.json')
if not os.path.exists(config_file):
    raise Exception("bert_config does not exist")
bert_config = BertModel_module.BertConfig.from_json_file(config_file)

max_seq_len = 256
num_labels = 9

def save():
    input_ids = tf.placeholder(tf.int32, (None, max_seq_len), 'input_ids')
    input_mask = tf.placeholder(tf.int32, (None, max_seq_len), 'input_mask')
    segment_ids = None
    #这里简单使用分类，具体根据自己需求修改
    probabilities = create_model(bert_config, False, input_ids, input_mask, segment_ids, num_labels, False)
    #根据自己的模型进行编写
    freeze_config = {
        "input_tensor": {
            'input_ids': input_ids,
            'input_mask': input_mask
        },
        "output_tensor": {
            "pred_ids": probabilities
        },
        "ckpt_filename": r"/home/build/fastertransformer/cmake_cuda1131/cmake38/tensorflow/tensorflow_bert/mrpc_output/model.ckpt-343",  # 训练ckpt权重
        "save_pb_file": r"/root/test.pb",# 保存pb文件
        "save_pb_path_serving": None,  # tf_serving 保存模型路径 ，None则不保存
    }
    ret = fasterseq.freeze(freeze_config)
    print(ret)

    # 查看
    fasterseq.pbshow(freeze_config['save_pb_file'])
save()


```