[Docker + Flask] 간단한 추론 API 서버 띄우기

Docker Container와 Flask를 활용하여 간단한 딥러닝 모델 추론 API 서버를 만들고자 한다.

추론서버 환경은 ubuntu를 적극 권장하며, 본 포스트에서 주요하게 다루고자 하는 부분은 다음 두 가지이다.

• 추론 API 서버 띄우기(Docker + Flask)
• request 라이브러리를 사용한 파일 전송

 

API 개념도

 

0. 간단한 Flask tutorial

가장 minimal한 Falsk application은 아래와 같으며, 해당 코드를 flask_test.py로 저장하자.

from flask import Flask

app = Flask(__name__)

@app.route('/hello')
def hello():
    return 'Hello World!'

Flask 공식문서에 따르면 위 코드에 대한 설명은 다음과 같다.

• Flask class를 import 해준다.
• Flask class의 instance를 생성해준다. 이때, 첫번재 argument는 application module이나 package의 이름을 넣는다. Flask가 template, static file을 찾을 수 있게 하는 값인데, 대부분의 경우 __name__을 넣어주면 된다.
• 여기가 중요한데, @app.route() decorator로 실행하고자 하는 함수를 감싸 해당 URL의 요청이 왔을 때, 해당하는 함수를 실행하여 리턴값을 줄 수 있게 한다. 위 예시에서는 /hello로 접속 시 "Hello World!"가 출력된다.

 

이제 terminal을 열어 아래 명령어로 flask_test.py를 실행시키고, 웹브라우져를 켜 접속해보면 hello world가 잘 나오는 것을 확인할 수 있다. 참고로 development는 개발 시 debugging 모드를 사용하는 옵션이다.

FLASK_ENV=development FLASK_APP=flask_test.py flask run

 

Hello World!

 

이 간단한 tutorial은 하나의 머신에서 서버를 띄우고 접속도 한 예시이다. 추론서버를 위해서는 remote 환경에서 접속이 가능하도록 해야한다.

 

 

1. REQUEST 부분

Pytorch 공식 문서를 참고하여 정리하고자 한다. 모델은 DenseNet 121이며, Pytorch에서 제공하는 공식 모델을 사용하겠다.

먼저 REQUEST를 통해 데이터를 Model Server에 전달하는 부분을 정리하고자 한다.

 

1) API 정의

API로 주고받는 데이터를 아래와 같이 정의해보자.

• request: image
• response: class_id, class_name(영어), answer(한글)

 

2) request 전달

API서버에 전달할 HTTP REQUEST를 정의하고 전달하는 과정이다. python requests 라이브러리를 사용하여 API 서버 주소인 192.168.xxx.xxx에 post로 이미지 파일을 전달할 것이다. 아래 코드를 client쪽에 send_data.py라는 이름으로 저장해두자.

import requests

def send_data(url):
    files = {
        'file':open('./path/to/image.jpg', 'rb')
    }
    res = requests.post(url,files=files)
    return res.text

url = "http://192.168.xxx.xxx:5000/predict"
print(send_post(url))

url = "http://192.168.xxx.xxx:5000/predict"인 이유는 Flask의 기본 포트가 5000이고, Model Server에 predict 함수를 만들어 추론 결과를 얻으려고 할 것이기 때문이다.

 

 

2. RESPONSE 부분

HTTP POST로 전달받은 데이터(이미지)를 모델에 입력하고 추론값을 반환하여 response를 생성하는 과정이다.

 

1) 이미지 preprocessing

아래와 같은 이미지가 들어오면 이를 Tensor로 변환해야 한다. 

zzz..

 

import io
import torchvision.transforms as transforms
from PIL import Image

def transform_image(image_bytes):
    my_transforms = transforms.Compose([transforms.Resize(255),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize(
                                            [0.485, 0.456, 0.406],
                                            [0.229, 0.224, 0.225])])
    image = Image.open(io.BytesIO(image_bytes))
    return my_transforms(image).unsqueeze(0)

 

jpeg 형식 이미지를 입력으로 받아야하며, 변환된 Tensor는 다음과 같다.

with open("../some/image.jpeg", 'rb') as f:
    image_bytes = f.read()
    tensor = transform_image(image_bytes=image_bytes)
    print(tensor)
    
>>>> tensor([[[[ 0.4508,  0.4166,  0.3994,  ..., -1.3473, -1.3302, -1.3473],
          	[ 0.5364,  0.4851,  0.4508,  ..., -1.2959, -1.3130, -1.3302],
          	[ 0.7077,  0.6392,  0.6049,  ..., -1.2959, -1.3302, -1.3644],
          	...,

 

2) 모델 추론 값 생성

이미지 preprosseccing + 모델 추론 결과를 얻는 과정이다. imagenet_class_index.json은 이미지 분류 결과의 클래스명을 담고 있는 데이터로 여기서 받을 수 있다.

from torchvision import models
import json

# 이미 학습된 가중치를 사용하기 위해 `pretrained` 에 `True` 값을 전달
model = models.densenet121(pretrained=True)

# 모델을 추론에만 사용할 것이므로, `eval` 모드로 변경
model.eval()

# ImageNet 분류 ID와 ImageNet 분류명의 쌍 정보
imagenet_class_index = json.load(open('../path/to/imagenet_class_index.json'))

def get_prediction(image_bytes):
    tensor = transform_image(image_bytes=image_bytes)  # i.에서 정의
    outputs = model.forward(tensor)
    _, y_hat = outputs.max(1)
    predicted_idx = str(y_hat.item())
    
    return imagenet_class_index[predicted_idx]

 

이미지를 샘플로 넣어보면 아래와 같이 예측값을 얻을 수 있다.

with open("../some/image.jpeg", 'rb') as f:
    image_bytes = f.read()
    print(get_prediction(image_bytes=image_bytes))
    
>>>> ['n02124075', 'Egyptian_cat']

 

3) response 생성

모델 추론값 response를 생성하여 전달하는 과정이다. 전달할 데이터는 class_id, class_name(엉어), answer(한글)이다.

핵심은 추론에 사용할 함수 predict()를 정의해주는 것과 Flask에서 make_response 모듈을 활용하는 것이다. url="http://192.168.xxx.xxx:5000/predict"를 통해 response를 반환하고, request 데이터는 post로 전달할 계획이기 때문에 @app.route('/predict', mothods=['POST']) decorator로 predict()를 감싸준다.

from flask import request
from flask import make_response

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
	    # 전달받은 request에서 이미지 데이터 받고 byte로 변환
        file = request.files['image']
        img_bytes = file.read()
        
        # 추론값 생성
        class_id, class_name = get_prediction(image_bytes=img_bytes)
        
        # response로 전달할 데이터 생성
        answer = "이 동물은 %s 입니다"%(class_name)
        
        # 주의!! #
        # jsonify({'class_id': class_id, 'class_name': class_name})와 같이 jsonify를 사용하면
        # json.dump()와 똑같이 ascii 인코딩을 사용하기 때문에 한글 깨짐
        
        # 이 방법으로 response 전달해야 함(!!!!)
        res = {
        	'class_id' : class_id,
            'class_name' : class_name,
            'answer' : answer
        }
        res = make_response(json.dumps(res, ensure_ascii=False))
        res.headers['Content-Type'] = 'application/json'
        
        return res

 

이제 1), 2), 3)을 하나의 파일에 다 합쳐 model_api.py로 저장하여 Model Server쪽에 두자.

model_api.py

0.00MB

 

 

3. API 서버 생성 및 테스트

1) Docker Container 기반 API 서버 띄우기

Model Server 192.168.xxx.xxx에 container로 추론 서버를 띄우는 과정이다. Pytorch 이미지를 사용할 것이며, Model Server에서 아래 코드로 container를 띄우고자 한다. 주요 argument, directory는 다음과 같다.

• image: pytorch/pytorch:1.8.1-cuda11.1-cudnn8-devel
• mount: source에 model_api.py가 저장되어 있는 directory 입력
• port: Flask가 사용하는 기본 port인 5000을 지정

 

먼저 아래 코드로 container을 생성한다.

# container 생성
sudo docker run -d -it --name (container 이름) --gpus "device=0" -p 5000:5000 --mount type=bind,source=/PATH/TO/model_api.py,target=/root pytorch/pytorch:1.8.1-cuda11.1-cudnn8-devel

만약 host machine에서 다른 port를 사용하고 싶다면 아래와 같이 사용하면 된다

# container 생성
sudo docker run -d -it --name (container 이름) --gpus "device=0" -p 12345:5000 ...

 

 

만든 container를 실행시킨 후 접속해서 Flask를 설치해주고

# container 실행 후 접속
sudo docker start (container 이름)
sudo docker attach (container 이름)

# container 접속 후 실행
pip install --upgrade pip && pip install flask

 

model_api.py를 실행시켜 둔다. 참고로 nohup으로 실행시키면 terminal 창을 꺼도 계속 실행되며, nohup 사용시 마지막에 &을 붙여주는게 좋다고 한다.

# API 추론 모델 실행
python model_api.py  		  # python process 실행 시
nohup python model_api.py &   # background에서 python process 실행 시

model_api.py 실행 화면

 

2) API 테스트

다 왔다! 이제 client 환경에서 데이터를 전달해보자. client 환경에서 send_data.py를 실행해주자.

python send_data.py

 

실행결과 client 환경에서는 추론 결과 response가 잘 출력되며, Model Server에서도 HTTP 200으로 결과가 잘 전달된 것을 확인할 수 있다.

 

 

 

다음 글들에서는 통해 한글 인코딩, requests 모듈에 대해 정리해보고자 한다.

긴 글 읽어주셔서 감사합니다.