tensorflow serving部署Bert预训练模型

灰太狼 2022-12-31 13:23 239阅读 0赞

**目前没有整理完善，先留个坑~**

--------------------

### Bert模型介绍 ###

BERT的关键技术创新是将Transformers双向训练作为一种流行的注意力模型应用到语言建模中。Masked LM (MLM)在向BERT输入单词序列之前，每个序列中有15%的单词被\[MASK\]token替换。然后，该模型试图根据序列中其他非MASK词提供的上下文来预测MASK词的原始值。

本文主要记录使用tensorflow serving部署训练好的bert模型，并根据模型获取句子向量表示。

### ckpt转saved\_model格式 ###

google bert原始预训练模型保存的事ckpt格式，用tfserving部署需要saved\_model的pb格式，这里需要一个转化过程。

import json
    import os
    import tensorflow as tf
    import argparse
    
    import modeling
    
    def create_model(bert_config, is_training, input_ids):
        # 通过传入的训练数据，进行representation
        model = modeling.BertModel(config=bert_config, is_training=is_training, input_ids=input_ids)
        output = model.get_pooled_output()
        # output = model.get_sequence_output()
    
        return output
    
    def transfer_saved_model(args):
    
        gpu_config = tf.ConfigProto()
        gpu_config.gpu_options.allow_growth = True
        sess = tf.Session(config=gpu_config)
    
        print("going to restore checkpoint")
        bert_config_file = os.path.join(args.model_path, 'bert_config.json')
        bert_config = modeling.BertConfig.from_json_file(bert_config_file)
    
        input_ids = tf.placeholder(tf.int32, [None, args.max_seq_len], name="input_ids")
        output = create_model(bert_config=bert_config, is_training=False, input_ids=input_ids)
    
        saver = tf.train.Saver()
        saver.restore(sess, tf.train.latest_checkpoint(args.model_path))
    
        tf.saved_model.simple_save(sess, args.export_path,
                                   inputs={ 'input_ids': input_ids},
                                   outputs={ "outputs": output})
        print('savedModel export finished.')
        
    if __name__ == '__main__':
        parser = argparse.ArgumentParser(description='Trans ckpt file to .pb file')
        
        parser.add_argument('-model_path', type=str, required=True, help='dir of a pretrained BERT model')
        parser.add_argument('-export_path', type=str, default=None, help='export model path')
        parser.add_argument('-max_seq_len', type=int, default=128, help='maximum length of a sequence')
        args = parser.parse_args()
    
        transfer_saved_model(args)

运行上述python脚本

python3 export_bert_test.py -model_path uncased_L-12_H-768_A-12 \ 
    -export_path bert_output -max_seq_len 128

### Docker部署模型 ###

使用docker版tensorflow serving部署模型。

nvidia-docker run -p 8525:8525 -p 8526:8526 --name=bert_model \
    --mount type=bind,source=$(pwd)/bert_output,target=/models/bert_output \
    -t --entrypoint=tensorflow_model_server tensorflow/serving:1.12.0-gpu \
    --port=8525 --rest_api_port=8526 \
    --enable_batching=true --file_system_poll_wait_seconds=300 \
    --grpc_channel_arguments=“grpc.max_connection_age_ms=5000” \
    --per_process_gpu_memory_fraction=0.4 \
    -- model_config_file=/models/bert_output/tfserving.conf

### http和grpc请求 ###

import tokenizer
    
    def process_input(texts, max_len):
        def tokenize_input(text):
            tokens = tokenizer.tokenize(text)
            text_ids = tokenizer.convert_tokens_to_ids(tokens)
            while len(text_ids) < max_len:
                text_ids.append(0)
            return text_ids
    
        input_ids = []
        for text in texts:
            text_ids = tokenize_input(text)
            input_ids.append(text_ids)
        
        return input_ids

**http请求**

import numpy as np
    import requests
    
    def predict_http(input_texts):
        SERVER_URL = "http://localhost:8526/v1/models/bert_output:predict"
        
        input_ids = process_input(input_texts, max_len=128)
        input_ids = np.array(input_ids).tolist()
        input_data = { "input_ids": input_ids}
        
        request_data = { "signature_name": "serving_default", "instances": input_data}
        request_data = json.dumps(request_data)
    
        response = requests.post(SERVER_URL, data=request_data)
        result = response.json()
        pred_value = result['outputs']
    
        return pred_value

**grpc请求**

import grpc
    import tensorflow as tf
    from tensorflow.core.framework import types_pb2
    from tensorflow_serving.apis import predict_pb2
    from tensorflow_serving.apis import prediction_service_pb2_grpc
    
    host = "localhost"
    port = "8526"
    
    # initialize grpc channel
    channel = grpc.insecure_channel('{host}:{port}'.format(host=host, port=port))
    
    def predict_grpc(input_texts, max_len=128):
    
        stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    
        result = []
    
        def predict_request(input_ids):
            """模型预测请求."""
            input_ids = np.array(input_ids)
            input_tensor = tf.make_tensor_proto(input_ids, shape=input_ids.shape, dtype=tf.int32)
            try:
                request = predict_pb2.PredictRequest()
                request.inputs["input_ids"].ParseFromString(input_tensor.SerializeToString())
                request.model_spec.name = "bert_output"
                request.model_spec.signature_name = "serving_default"
                response = stub.Predict(request, 50)
                _result = tf.make_ndarray(response.outputs["outputs"]).tolist()
                result.extend(_result)
            except Exception as e:
                print(e)
    
            input_ids = process_input(input_texts, max_len)
            predict_request(input_ids)
    
        return result