Skip to content

Use BERT to train a classification model and deploy the model by tensorflow serving

Notifications You must be signed in to change notification settings

nghuyong/bert-classification-tf-serving

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

bert-classification-tf-serving

Introduction

Use BERT to train a classification model and deploy the model by tensorflow serving. Then we can use REST API to do online prediction.

Get Started

The whole experiment is based on Chnsenticorp data set, which is a two-class data set of positive and negative emotions.

0. Prepare Pre-train model

Download the Chinese bert model chinese_L-12_H-768_A-12, then unzip and move to the models dir.

1. Train a model

sh fine-tuning.sh

2. Do prediction and export model

We need to change the checkpoint to the tensorflow serving's format.

sh export-model.sh

Then the structure of export-model will be:

.
└── 1569141360
    ├── saved_model.pb
    └── variables
        ├── variables.data-00000-of-00001
        └── variables.index

3. Deploy model

We use tensorflow serving docker version to deploy the model on CPU OR GPU.

a. On CPU

docker run -p 8500:8500 -p 8501:8501 --mount type=bind,source=$(pwd)/models/export-model,target=/models/chnsenticorp -e MODEL_NAME=chnsenticorp -t tensorflow/serving

b. On GPU

You need to install nvidia-docker environment first!

docker run --runtime=nvidia -e CUDA_VISIBLE_DEVICES=0 -p 8500:8500 -p 8501:8501 --mount type=bind,source=$(pwd)/models/export-model,target=/models/chnsenticorp -e MODEL_NAME=chnsenticorp -t tensorflow/serving:1.12.0-gpu

4. Predict online

export PYTHONPATH=$(pwd)/google-bert
python predict_online.py

a. On CPU

Input test sentence:
跟住招待所没什么太大区别。 绝对不会再住第2次的酒店!
negative pro:0.96344769 positive pro:0.0365523435 time consuming:198ms
Input test sentence:
LED屏就是爽,基本硬件配置都很均衡,镜面考漆不错,小黑,我喜欢。
negative pro:0.0291572567 positive pro:0.970842779 time consuming:195ms

b. On GPU

Input test sentence:
跟住招待所没什么太大区别。 绝对不会再住第2次的酒店!
negative pro:0.963448 positive pro:0.0365524 time consuming:92ms
Input test sentence:
LED屏就是爽,基本硬件配置都很均衡,镜面考漆不错,小黑,我喜欢。
negative pro:0.0291573 positive pro:0.970843 time consuming:87ms

The probability on GPU and GPU is the same, but the speed on GPU is faster 50%+ than on CPU!

The max token length is 128 and the CPU and GPU configuration is as follows:

CPU GPU
Intel(R) Xeon(R) CPU E5-2650 v4 @ 2.20GHz Tesla V100

Conclusion

This is a demo project for classification by BERT and deploy the model by tensorflow serving.

We use tensorflow serving docker version, so we can use docker stack technology to easily expand for load balance in production environment.

About

Use BERT to train a classification model and deploy the model by tensorflow serving

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published