파이토치 모델 서빙하고 내보내기 (도커, 토치서브, 토치스크립트, ONNX 활용)

훈련이 완료된 ML/DL 모델은 실제 서비스에 적용하기 위해 프로덕션 레벨로 배포하고 서빙해야한다. 이를 위한 모델 서빙 종류는 로컬에 플라스크로 모델 서버 구축, 도커로 모델 서빙 환경 구축, 토치서브로 모델 서빙하는 방법, 클라우드로 모델 서빙하는 방법이 있다. 모델 서빙은 모델 훈련이 완료되고 저장되어 있어 재사용가능 할때 이루어지므로 간단히 모델 저장 방법에 대해 언급하고 넘어가자.

 

파이토치 모델 저장 및 로딩 방법

파이토처리 모델을 저장하고 로딩하는 방법은 2가지가 있다. 첫 번째 방법은 가장 간단한 방식으로 전체 모델 객체를 저장하고 로딩하는 방식이다.

torch.save(model, PATH)
model = torch.load(PATH)

모델 전체를 저장하므로 간편하다. 하지만 이 방식은 경우에 따라 문제가 발생할 수 있다. 그 이유는 모델 객체 전체를 저장하기 때문에 모델 객체 내부에 매개변수, 모델 클래스, 디렉터리 구조까지 함께 저장한다. 따라서 만약 추후 디렉터리 구조라도 변경된다면 모델 로딩에 실패하게 되고 문제 해결이 어려워질 수 있다. 따라서 모델 매개변수만 저장하는 아래 두 번째 방법을 사용하는 것이 좋다.

torch.save(model.state_dict(), PATH)
model = ConvNet()
model.load_state_dict(torch.load(PATH))

모델 매개변수만 저장하기 위해 state_dict() 메서드를 사용하고, 추후 빈 모델 객체를 인스턴스화해 빈 모델 객체에 매개 변수를 로딩해 사용하는 방법이다. 간단히 파이토치 모델 저장하는 방법을 알아보았으니 다음으로 이 모델을 배포하여 서빙하는 방법에 대해 알아보자. 

 

1. 플라스크로 로컬에 모델 서버 구축하기

가장 간단하게 설치해 사용할 수 있는 플라스크를 서버로 사용한다. 플라스크 서버에 입력(추론 요청)이 들어오면 추론한 결과를 출력값으로 하여 되돌려 보내준다. 중요한 것은 플라스크 서버 내부에서 추론해 결과를 돌려주는 일이므로 추론 함수 작성이 필요하다. 예시를 위해 MNIST 숫자를 예측하는 모델이라 가정한다. 그렇다면 다음과 같은 추론 함수를 작성할 수 있다. 

def run_model(input_tensor):
    model_input = input_tensor.unsqueeze(0)
    with torch.no_grad():
        model_output = model(model_input)[0]
    model_prediction = model_output.detach().numpy().argmax()
    return model_prediction

 

이 추론 함수를 플라스크 서버 내부에서 동작시키려면 플라스크 서버의 코드는 다음과 같은 형태로 작성할 수 있다. 

 

import os
import json
import numpy as np
from flask import Flask, request

import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.cn1 = nn.Conv2d(1, 16, 3, 1)
        self.cn2 = nn.Conv2d(16, 32, 3, 1)
        self.dp1 = nn.Dropout2d(0.10)
        self.dp2 = nn.Dropout2d(0.25)
        self.fc1 = nn.Linear(4608, 64) # 4608 is basically 12 X 12 X 32
        self.fc2 = nn.Linear(64, 10)
 
    def forward(self, x):
        x = self.cn1(x)
        x = F.relu(x)
        x = self.cn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dp1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dp2(x)
        x = self.fc2(x)
        op = F.log_softmax(x, dim=1)
        return op
    
model = ConvNet()
PATH_TO_MODEL = "./convnet.pth"
model.load_state_dict(torch.load(PATH_TO_MODEL, map_location="cpu"))
model.eval()

def run_model(input_tensor):
    model_input = input_tensor.unsqueeze(0)
    with torch.no_grad():
        model_output = model(model_input)[0]
    model_prediction = model_output.detach().numpy().argmax()
    return model_prediction

def post_process(output):
    return str(output)

app = Flask(__name__)

@app.route("/test", methods=["POST"])
def test():
    data = request.files['data'].read()
    md = json.load(request.files['metadata'])
    input_array = np.frombuffer(data, dtype=np.float32)
    input_image_tensor = torch.from_numpy(input_array).view(md["dims"])
    output = run_model(input_image_tensor)
    final_output = post_process(output)
    return final_output

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8890)

 

위와 같이 작성된 플라스크 서버를 구동시키면 입력을 기다릴 것이고 입력 request를 날려주는 함수를 하나 작성하면 다음과 같이 작성할 수 있다.

 

import io
import json
import requests
from PIL import Image
from torchvision import transforms

def image_to_tensor(image):
    gray_image = transforms.functional.to_grayscale(image)
    resized_image = transforms.functional.resize(gray_image, (28, 28))
    input_image_tensor = transforms.functional.to_tensor(resized_image)
    input_image_tensor_norm = transforms.functional.normalize(input_image_tensor, (0.1302,), (0.3069,))
    return input_image_tensor_norm

image = Image.open("./digit_image.jpg")
image_tensor = image_to_tensor(image)
dimensions = io.StringIO(json.dumps({'dims': list(image_tensor.shape)}))
data = io.BytesIO(bytearray(image_tensor.numpy()))

res = requests.post('http://localhost:8890/test', files={'metadata': dimensions, 'data' : data})
response = json.loads(res.content)

print("Predicted digit :", response)

 

위와 같이 작성한 추론 요청 함수를 실행시키면 플라스크 모델 서버로부터 추론한 결과를 받을 수 있게 된다. 하지만 이와 같이 구축된 모델 추론 파이프라인을 다른 환경에서 동일하게 구축하려면 수동으로 동일한 라이브러리 설치부터 폴더구조를 맞춰주고 파일을 복사하는 등의 작업이 필요하다. 때문에 이 환경을 그대로 복제하여 다른 환경에서도 사용할 수 있도록 하는 확장성이 필요하다. 이 때 사용할 수 있는 것이 도커다. 도커를 통해 손쉽게 위 환경을 복제하고 실행하는 방법에 대해 알아보자. 

 

 

2. 도커로 모델 서빙 환경 구축하기

도커를 사용하면 서버 구축에 사용되었던 소프트웨어 환경을 손쉽게 만들 수 있다.

 

2.1 Dockerfile 만들기

도커를 사용하기 위해서는 가장 먼저 Dockerfile을 만들어야 한다. Dockerfile을 실행하면 결과적으로 위 플라스크로 구축했던 모델 서빙환경이 그대로 재현되어야 한다. Dockerfile 실행하면 도커 이미지가 생성되고 이 과정을 이미지 빌드라고 한다. 이를 위해 Dockerfile에 다음과 같은 스크립트를 작성할 수 있다.

FROM python:3.8-slim 

RUN apt-get -q update && apt-get -q install -y wget

COPY ./server.py ./
COPY ./requirements.txt ./

RUN wget -q https://raw.githubusercontent.com/wikibook/mpytc/main/Chapter10/convnet.pth
RUN wget -q https://github.com/wikibook/mpytc/raw/main/Chapter10/digit_image.jpg

RUN pip install --no-cache-dir -r requirements.txt

USER root
ENTRYPOINT ["python", "server.py"]

 

간단히 각 명령에 대해 설명하자면 FROM을 통해 도커에 python3.8이 포함된 표준 리눅스 OS를 가져오도록 지시하고, RUN을 통해 업데이트 하고 wget을 다운로드 한다. 이후 COPY를 통해 로컬에 만들어두었던 서버 파일과 환경구축에 필요한 requirements.txt 파일을 복사해준다. 이후 추론에 필요로한 모델과 예제 이미지를 다운로드 받고, pip install을 통해 파이썬 라이브러리를 모두 설치해준다. 참고로 requirements.txt에는 다음과 같은 라이브러리가 들어있다. (만약 동작하지 않는다면 Flask를 업데이트 할 것) 

torch==1.5.0
torchvision==0.5.0
Pillow==6.2.2
Flask==1.1.1

USER를 통해 루트 권한을 부여하고 ENTRYPOINT를 통해 python server.py 명령이 실행되면 도커상에서 플라스크 모델 서버가 실행된다. 

 

2.2 도커 이미지 빌드

도커 파일이 작성되었으면 이를 빌드해주어 도커 이미지로 만들어주어야 한다. 이를 위해 다음과 같은 명령어를 사용한다.

docker build -t <tag name> .

여기서 <tag name>은 임의 설정이다. 

 

2.3 도커 이미지 배포 (실행)

도커 이미지 빌드를 통해 만들어진 도커 이미지를 다음과 같은 명령 실행을 통해 배포가 가능하다.

docker run -p 8890:8890 <tag name>

 

위 명령이 실행되면 다음과 같이 도커 이미지가 실행되어 모델 서빙이 가능하도록 플라스크 서버가 구동되는 것을 확인할 수 있다.

 

구동된 플라스크 서버에 다시 추론을 요청하는 request를 날리면 이에 대한 응답이 결과로 반환되는 것을 확인할 수 있다.

 

도커 실행을 중지하려면 현재 실행 중인 도커 컨테이너를 확인해야 하며 이는 docker ps -a 명령으로 확인할 수 있다. 그러면 CONTAINER ID를 확인할 수 있고 이를 복사하여 docker stop <CONTAINER ID> -t 0을 실행시켜주면 도커 컨테이너가 중지된다. 만약 도커 컨테이너를 삭제하고 싶다면 docke rm <CONTAINER ID> 명령을 실행하면 된다. 빌드했던 도커 이미지까지 삭제하려면 docker images 명령을 통해 <tag name>을 확인하고 docker rmi <tag name>을 통해 삭제할 수 있다.

 

3. 토치서브로 모델 서빙하기

torchserve는 파이토치 모델 서버 라이브러리다. 메타와 AWS에서 만들었고, 파이토치 모델 배포를 도와주는 역할을 한다. torchserve를 사용하기 위해서 Java 11 SDK 설치와, pip install torchserve torch-model-archiver 명령 실행을 통해 라이브러리 설치가 우선적으로 필요하다. 여기서 torch-model-archiver는 압축 라이브러리다. 입력값을 3개 받아 .mar 파일로 만들어주는 역할을 한다. 입력값 3개란 1. 모델 클래스 파일 2. 훈련된 모델 파일 3. 핸들러(전처리, 후처리) 파일이다.

 

1. 클래스 파일은 모델 레이어가 구성된 파일로, 앞서 플라스크 서버 구축때 사용할 때 정의했던 ConvNet 클래스를 별도 파일로 만들어준 것이다. 

2. 훈련된 모델 파일은 convnet.pth와 같이 모델 학습이 완료된 파일이며

3. 핸들러는 torchvision의 transform 클래스와 같이 전처리를 할 수 있는 로직이나 별도의 후처리 로직이 담긴 파일이다.

 

torch-model-archiver를 사용해 아래 명령을 실행시켜주면 목표로하는 convnet.mar 파일이 생성된다. 

torch-model-archiver --model-name convnet --version 1.0 \
--model-file ./convnet.py --serialized-file ./convnet.pth --handler ./convnet_handler.py

이후 아래 명령을 통해 새 디렉터리를 생성하고 생성한 convnet.mar 파일을 옮겨준다.

mkdir model_store
mv convnet.mar model_store

이후 토치서브를 이용해 모델 서버를 런칭하는 명령 실행을 통해 서버를 구동해준다.

torchserve --start --ncs --model-store model_store --models convnet.mar

이후 curl을 활용해 토치서브 모델 서버에 추론 요청을 수행할 수 있다. 핸들러에 의해 어떤 입력이라도 전처리되어 텐서로 바뀌기 대문에 별도의 코드를 작성하지 않아도 된다.

curl http://127.0.0.1:8080/predictions/convnet -T ./digit_image.jpg

결과적으로 추론한 결과 숫자가 터미널에 출력될 것이다. 만약 토치서브로 구동한 서버를 종료하고 싶다면 torchserve --stop 명령을 통해 중지할 수 있다. 덧붙여 만약 서버의 모델 서빙 여부를 확인하기 위한 핑을 보내는 명령은 다음과 같이 사용할 수 있다. 참고로 포트 정보는 torchserve 실행에 의해 서버가 실행될 때 로그를 통해 확인 가능하다.

curl http://localhost:8081/models

 

 

4. 토치스크립트로 범용 파이토치 모델 만들기

지금까지는 로컬 플라스크 서버, 도커 환경, 토치서브 환경을 통해 모델 서버를 구현했다. 하지만 이는 파이썬 스크립트 환경에서 이루어진 것이다. 훈련한 모델이 반드시 훈련했던 환경에서 모델 서빙이 이뤄진다 할 수 없다. 파이썬이 실행되지 않는 외부환경에서도 실행될 수 있다. 따라서 파이토치 모델을 C++과 같은 타 언어에서도 실행될 수 있도록 중간 표현으로 만들 필요성이 있다. 이 때 토치스크립트를 사용할 수 있다. 토치스크립트는 별도의 라이브러리는 아니고 파이토치 내부에서 모델 연산 최적화를 위해 사용하는 JIT compiler에서 사용된다. 토치스크립트는 간단히 torch.jit.script(model) 또는 torch.jit.trace(model, input)를 실행시켜 만들 수 있다. 

 

파이토치 모델을 토치스크립트로 컴파일하기 위해선, 언급했듯 두 가지 방식이 존재한다. 첫 번째는 trace 방식이고 두 번째는 script 방식이다. trace 방식을 위해선 모델과 더미(임의) 입력값이 필요하다. 더미 입력값을 모델에 넣어서 입력값이 어떻게 흐르는지 추적해 기록한다. trace 방식을 사용해 모델을 만든다면 다음과 같은 형태로 만들 수 있다.

model = ConvNet()
PATH_TO_MODEL = "./convnet.pth"
model.load_state_dict(torch.load(PATH_TO_MODEL, map_location="cpu"))
model.eval()
for p in model.parameters():
    p.requires_grad_(False)
demo_input = torch.ones(1, 1, 28, 28)
traced_model = torch.jit.trace(model, demo_input)
# print (traced_model.graph)
# print(traced_model.code)
torch.jit.save(traced_model, 'traced_convnet.pt')
loaded_traced_model = torch.jit.load('traced_convnet.pt')

중간에 있는 torch.jit.trace() 메서드를 통해 토치스크립트 형식의 객체를 만들어주는 것이 핵심이다. 이후 모델을 저장하게 되면 C++과 같은 다른 언어에서도 파이토치 모델을 로딩해 추론할 수 있다. 위를 살펴보면 model과 loaded_trace_model 두 개가 있는데 향후 모델에 입력값을 넣고 추론해보면 추론 결과는 완벽히 동일하게 출력된다.

 

하지만 이러한 trace 방식을 이용하면 한 가지 큰 문제가 발생할 수 있다. 예를 들어 모델 순전파가 if나 for문과 같은 제어 흐름으로 구성된다면 trace는 여러 가능한 경로 중 하나만 토치 스크립트로 렌더링할 것이다. 따라서 기존 모델과 동일성을 보장할 수 없다. 이를 해결하기 위해서는 script 방식을 사용해 토치 스크립트로 컴파일 해야 한다. script 방식도 위 trace 방식과 동일하다. 다만 차이점이 있다면 더미 입력이 사용되지 않는 것이 특징이다. 

model = ConvNet()
PATH_TO_MODEL = "./convnet.pth"
model.load_state_dict(torch.load(PATH_TO_MODEL, map_location="cpu"))
model.eval()
for p in model.parameters():
    p.requires_grad_(False)
scripted_model = torch.jit.script(model)
# print(scripted_model.graph)
# print(scripted_model.code)
torch.jit.save(scripted_model, 'scripted_convnet.pt')
loaded_scripted_model = torch.jit.load('scripted_convnet.pt')

마찬가지로 중간에 있는 torch.jit.script() 메서드를 통해 토치스크립트 형식의 객체를 만들어주는 것이 핵심이다. 또 model과 loaded_scripted_model에 모두 입력값을 넣고 추론해보면 추론 결과가 동일하게 출력된다. 이렇게 script 방식을 사용해 토치스크립트 코드로 만들면 trace 방식에 비해 정확성을 더 얻을 수 있다. 다만 script 방식이 가지는 단점이 있다면 파이토치 모델이 토치스크립트에서 지원하지 않는 기능을 포함하면 동작할 수 없다. 이 때는 모델 순전파에 if for문 로직을 제거한 뒤 trace 방식을 사용해야 한다.

 

결론을 이야기하자면 사실상 토치스크립트는 모델 서빙을 위해 필수적으로 사용해야 한다. 그 이유는 파이썬 내부에서는 전역 인터프리터 잠금(GIL)이 설정되어 있어 한 번에 한 쓰레드만 실행될 수 있어 연산 병렬화가 불가능하기 때문이다. 하지만 토치스크립트를 통해 중간 표현으로 바꾸어 범용 형식으로 만들어두면 연산 병렬화가 가능해져 모델 서빙 속도 향상이 가능해진다. 

 

5. ONNX로 범용 파이토치 모델 만들기

토치스크립트와 마찬가지로 ONNX 프레임워크를 사용해도 파이토치 모델을 범용화할 수 있다. 그렇다면 토치스크립트와의 차이점은 무엇일까? 토치스크립트는 파이토치에서만 사용할 수 있다면 ONNX는 텐서플로우나 이외의 딥러닝 라이브러리와 같이 더 넓은 범위에서 표준화 시킬 수 있다. 예를 들어 파이토치로 만든 모델을 텐서플로우에서 로드하여 사용할 수 있다. 이를 위해 ONNX 라이브러리 설치가 필요하다. pip install onnx onnx-tf 명령을 통해 필요한 라이브러리를 설치할 수 있다. 

 

학습한 파이토치 모델을 ONNX 포맷으로 바꾸어 저장하기 위해서는 다음과 같은 예시 코드를 사용할 수 있다.

model = ConvNet()
PATH_TO_MODEL = "./convnet.pth"
model.load_state_dict(torch.load(PATH_TO_MODEL, map_location="cpu"))
model.eval()
for p in model.parameters():
    p.requires_grad_(False)
demo_input = torch.ones(1, 1, 28, 28)
torch.onnx.export(model, demo_input, "convnet.onnx")

토치스크립트에서 trace 방식과 마찬가지로 더미 입력값이 필요하다. 핵심은 torch.onnx.export() 메서드를 사용해 onnx 포맷 형식으로 파일을 생성하는 것이다. 이후 텐서플로우 모델로 변환하기 위해서는 다음과 같은 예시 코드를 사용할 수 있다.

 

import onnx
from onnx_tf.backend import prepare

model_onnx = onnx.load("./convnet.onnx")
tf_rep = prepare(model_onnx)
tf_rep.export_graph("./convnet.pb")

 

tf.rep.export_graph() 메서드가 실행되면 텐서플로우에서 사용 가능한 모델 파일인 convnet.pb가 생성된다.

 

이후 텐서플로우에서 사용하기 위해 모델 그래프를 파싱하는 과정이 필요하다. (참고로 TF 1.5 버전) 

with tf.gfile.GFile("./convnet.pb", "rb") as f:
    graph_definition = tf.GraphDef()
    graph_definition.ParseFromString(f.read())
    
with tf.Graph().as_default() as model_graph:
    tf.import_graph_def(graph_definition, name="")
    
for op in model_graph.get_operations():
    print(op.values())

 

그래프를 파싱하고 그 결과를 출력하면 아래와 같이(간략화된) 그래프의 입출력 노드 정보를 확인할 수 있다. 

(<tf.Tensor 'Const:0' shape=(16,) dtype=float32>,)
...
(<tf.Tensor 'input.1:0' shape=(1, 1, 28, 28) dtype=float32>,)
...
(<tf.Tensor '18:0' shape=(1, 10) dtype=float32>,)

 

입출력 노드 정보를 기반으로 모델의 입출력을 지정해줄 수 있고 텐서플로우를 통해 추론하면 아래와 같이 예측 확률분포가 출력되는 것을 확인할 수 있다. 

model_output = model_graph.get_tensor_by_name('18:0')
model_input = model_graph.get_tensor_by_name('input.1:0')

sess = tf.Session(graph=model_graph)
output = sess.run(model_output, feed_dict={model_input: input_tensor.unsqueeze(0)})
print(output)

> [[-9.35050774e+00 -1.20893326e+01 -2.23922171e-03 -8.92477798e+00
  -9.81972313e+00 -1.33498535e+01 -9.04598618e+00 -1.44924192e+01
  -6.30233145e+00 -1.22827682e+01]]

 

지금까지의 전체 과정을 요약하자면 파이토치 모델을 서빙하는 방법과 모델을 범용화 시키는 방법에 대해 알아봤다. 모델 서빙을 위해서는 크게 3가지로 로컬에서 플라스크 서버를 실행시켜 서빙하는 방법, 플라스크 서버를 도커 이미지화 하는 방법, 토치서브를 통해 서빙하는 방법이 있었다. 도커 이미지화를 위해선 Dockerfile을 생성해주고 몇몇 명령어를 실행시켜주었고, 토치서브를 실행시키기 위해 torch-model-archiver를 통해 .mar 파일로 만들어준 뒤, 서버를 실행시켜 추론 결과를 응답해주었다. 다음으로 파이토치 모델의 범용화를 위해 토치스크립트를 사용해 모델 범용화하는 방법과 ONNX를 통해 모델 범용화하는 방법이다.

 

잘 이뤄지는 모델 학습이  중요한 만큼 서비스와 제품화를 위해 모델을 배포하고 서빙하는 과정도 중요하다. 특히 실제 사용되는 서비스에서는 속도가 중요하므로 토치스크립트나 ONNX를 통해 모델을 범용 포맷으로 만듦으로써 파이썬 속도 한계의 원인인 전역 인터프리터 잠금(GIL)으로부터 벗어날 수 있다. 이후 플라스크 서버나 토치서브 환경을 도커 컨테이너로 만들어 둔다면 개발 속도 측면이나 효율성 등의 측면에서 많은 이점을 얻을 수 있을 것이다. 글을 작성하며 사용한 예시 코드의 전체 코드는 아래 링크를 참조하였다.

 

코드: https://github.com/wikibook/mpytc/tree/main/Chapter10

 

Reference

[1] 실전! 파이토치 딥러닝 프로젝트