Tensorflow Serving vs Flask REST API

 

Tensorflow serving 을 이용하기 전에는 flask 를 이용하여 REST API 를 만들어 microservice를 진행하는 것을 당연하게 생각했었습니다. python 의 overhead  줄이기 위해 C++/golang 등 다른 framework 를 써야 된다는 것을 알았지만, 큰 영향이 없겠다 싶어 고민하지 않았던 것 같습니다.

 

"TensorFlow Serving is a flexible, high-performance serving system for machine learning models, designed for production environments."   

https://github.com/tensorflow/serving

 

 Tensorflow serving은 유연하고, 고성능 서빙 시스템으로 프로덕션 환경에 맞게 디자인 되었다고 되어있는데 그럼 flask REST API 와 속도차이가 얼마나 나는지 궁금해서 비교해 보았습니다.

 

 우선 Tensorflow Serving 경우 REST API 방식과 gRPC 방식 2가지를 모두 지원하고 있습니다. (gRPC 내용은 추가 필요)

1. Flask REST API 

2. REST API [ Tensorflow Serving ]

3. gRPC API [Tensorflow Serving]

 

# 우선 심플한 CNN toy 모델을 학습해서 모델을 저장해놓고 이를 비교 검증을 진행 하도록 하겠습니다.

간단히 Colab 을 통해서 학습을 진행하였습니다. 해당 모델 파일을 다운로드 받아 압축을 푼 후 적당한 위치에 저장합니다.

(아래는 Colab 실제 코드입니다.)

tf_model_mnist.tar.gz

0.55MB

## Toy Model CNN

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train/255.0, x_test/255.0

x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

class CNN(tf.keras.Model):
    def __init__(self):
        super(CNN, self).__init__()
        # 28,28,1
        self.layer1 = tf.keras.Sequential(
            [layers.Conv2D(32, kernel_size=3, 
                           padding='same',activation='relu'),
             layers.MaxPool2D(pool_size=(2,2), strides=2, padding='same')])
        # 14,14,32
        self.layer2 = tf.keras.Sequential(
            [layers.Conv2D(64, kernel_size=3, 
                           padding='same',activation='relu'),
             layers.MaxPool2D(pool_size=(2,2), strides=2, padding='same')])
        # 7,7,64
        self.flatten = layers.Flatten()
        self.fc = layers.Dense(10)

    def call(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.flatten(out)
        return self.fc(out)
        
## * from_logits=True - softmax 수행하기 전의 값을 사용 ##
criterion = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()
tf_model = CNN()
tf_model.compile(loss=criterion, optimizer=optimizer)

train_epochs =20
batch_size = 1024
model_path = 'tf_cnn_model'
version = '1'
save_path = f'{model_path}/{version}'

tf_model.fit(x_train, y_train, batch_size=batch_size, epochs=train_epochs)

pred_y = tf_model.predict(x_test, batch_size=batch_size)
accuracy = np.sum(np.argmax(pred_y, axis=1) == y_test)/len(y_test)
print(f'accuracy : {accuracy:>.4f}')

tf.keras.models.save_model(tf_model, save_path)
#!tar -zcvf tf_cnn_model.tar.gz tf_cnn_model

 

 다음 해야 할일은 tenserflow serving 을 실행할 차례인데, docker 로 이미지를 다운로드 받아  docker 실행 시 몇가지 설정을 추가로 해주셔야 합니다.

 

기본적인 docker 사용법 및 옵션은 아래 사이트를 보시면 친절하고 자세히 설명되어 있습니다.

초보를 위한 도커 안내서 - 설치하고 컨테이너 실행하기

docker run [OPTIONS] IMAGE[:TAG|@DIGEST] [COMMAND] [ARG...]

다음은 자주 사용하는 옵션들입니다.

옵션설명

-d	detached mode 흔히 말하는 백그라운드 모드
-p	호스트와 컨테이너의 포트를 연결 (포워딩)
-v	호스트와 컨테이너의 디렉토리를 연결 (마운트)
-e	컨테이너 내에서 사용할 환경변수 설정
–name	컨테이너 이름 설정
–rm	프로세스 종료시 컨테이너 자동 제거
-it	-i와 -t를 동시에 사용한 것으로 터미널 입력을 위한 옵션
–link	컨테이너 연결 [컨테이너명:별칭]

 

# Download the TensorFlow Serving Docker image and repo
docker pull tensorflow/serving

# Run Tensorflow Serving(docker)
docker run -t --rm -p 8500:8500 -p 8501:8501 \
    -v "/home/roadcom/workspace/models/tf_cnn_model:/models/mnist_model" \
    -e MODEL_NAME=mnist_model \
    tensorflow/serving &

 

결과:  gRPC(TF) > REST(TF) >> Flask > Flask+gunicorn 순으로 나타납니다.

Tenserflwo  Serving이 왜 중요한지를 보여주는 차이입니다.

실험 방법은 mnist test set 을 동일하게 로딩 후 요청 - 응답 간 10 회 반복 값입니다. 오래된 로컬 컴퓨터라 속도가 빠르지는 않지만, 표준편차가 낮은 것을 보아 비교에는 큰 문제가 없을 것같습니다.

Metric	gRPC (TF)	REST (TF)	Flask	Flask+gunicorn
AVG (s)	0.983	1.624	9.088	9.394
STD (s)	0.011	0.021	0.378	0.348

 

해당 코드는 아래 github 에 공유드립니다.

https://github.com/elentail/Serving.git

 

 

참조 :

초보를 위한 도커 안내서 - 설치하고 컨테이너 실행하기

(https://subicura.com/2017/01/19/docker-guide-for-beginners-2.html)