PyTorch CNN 모델 Fast API로 배포하기

여지껏 CNN 모델을 사용해 볼 기회는 많았는데, 모델의 결과를 fast API로 배포해 본 적은 없었다. 최근에 좋은 기회가 생겨서 배포를 시도해 봤는데 그냥 두면 곧 까먹을 것 같아서 기록한다. (주먹구구 식으로 일단 배포만 되게 만든거라 다양한 기능이 있거나 효율이 좋은 코드는 아니다.) 

 

+ fast API를 설정하는 방법에 대해서는 따로 적지 않았다. 나는 Pycharm 2022버전을 사용했다. 2023 버전에는 fast API가 자동으로 설정되어 있다고 한다. 근데 난 자꾸 오류떠서 그냥 2022 버전에 fast API + uvicorn 설정했다.

 

 

아래 코드를 만들면서 가장 도움이 된 레퍼런스는 https://tutorials.pytorch.kr/intermediate/flask_rest_api_tutorial.html이다. 파이토치 공식문서 최고~ ^^

 

라이브러리 불러오기

from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from pydantic import BaseModel
from PIL import Image, ImageFile
from typing import List
from mangum import Mangum
from io import BytesIO
import torchvision.models as models
import torchvision.transforms as transforms
import torch
import base64
import io

 

 

환경설정

app = FastAPI()
handler = Mangum(app)
ImageFile.LOAD_TRUNCATED_IMAGES = True

mangum은 AWS 서버에 배포하기 위해 설정해준 것이고 마지막 TRUNCATED는 이미지 파일 입력받을 때 오류가 자꾸 떠서 한 줄 넣어줬다. 오류가 안 뜬다면 굳이 설정할 필요 없다. 맨 윗줄만 고고링

 

 

 

# 모델 로드
model1 = models.resnet50(num_classes=3)
model2 = models.efficientnet_b0(num_classes=3)


# 모델 가중치 파일 로드
device = torch.device('cpu')

model1.load_state_dict(torch.load('~~.pth', map_location=device))
model1.eval()

model2.load_state_dict(torch.load('~~.pth', map_location=device))
model2.eval()

서버를 켤 때마다 모델을 학습시키는 것은 현실적으로 불가능 (난 GPU가 없는 허접따리이므로 CNN 모델 학습이 애초에 불가능하다.) 따라서 코랩에서 미리 학습시킨 가중치 파일들을 위 코드에서 불러온다. 모델 학습에는 GPU가 필요하지만 evaluation은 CPU로 충분히 가능하니 문제없다.

 

 

 

 

def preprocess_image1(image):
    transform = transforms.Compose([
        transforms.Resize((300, 300)),
        transforms.RandomResizedCrop(256),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)


def preprocess_image2(image):
    transform = transforms.Compose([
        transforms.Resize((370, 370)),
        transforms.CenterCrop(260), 
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)

각 모델에 이미지가 입력됐을 때 수행할 transform을 지정한다. training 과정과는 달리 해괴한 증강기법을 넣지 않아도 된다. 간단하게 resize, crop, normalize 정도만 하는것이 일반적인 것 같다. 

 

 

 

 

 

# 이미지 업로드와 추론을 처리하는 엔드포인트
@app.post("/predict")
async def predict(image_base64: str = Form(...)):
    image = Image.open(BytesIO(base64.b64decode(image_base64))).convert('RGB')


    preprocessed_image1 = preprocess_image1(image)
    with torch.no_grad():
        outputs1 = model1(preprocessed_image1)

    probabilities1 = torch.nn.functional.softmax(outputs1[0], dim=0)
    inner = probabilities1 * torch.tensor([1.0, 0.0, 0.5])
    risk = torch.sum(inner)
    predicted_class1 = torch.argmax(probabilities1).item()

파이토치 공식문서에서는 image file을 바로 입력받았지만, 나는 base64 형태로 받아온 후 이를 다시 Image로 변환하는 과정을 거쳤다. conver('RGB')를 수행하지 않으면 channel=4로 인식돼 에러가 발생한다. 자칫 넘어가면 간과하기 쉬운 에러라 이거 찾는 데 20분정도 쓴 것 같다 ㅋ

 

아래 inner (inner product) 과정은 내가 가진 task에서 수행했어야 하는 작업이므로, 다른 task에서는 무시해도 된다. 해당 task는 이미지의 위험도(risk)를 판단하는 것이었다. softmax를 이용해서 구한 세 개의 클래스의 확률 값에 [1.0, 0.0, 0.5]를 내적하면 위험도가 구해지는 과정이다. (데이터가 어떻게 구성되어있는지 기록하지 않아서 당연히 읽는 사람들은 모를 것 같지만.... 모델을 배포하는데는 중요하지 않으므로 그냥 skip해도 ㅇㅋ~)

 

 

 

    if class_names[predicted_class1] != 'good':
        preprocessed_image2 = preprocess_image2(image)

        with torch.no_grad():
            outputs2 = model2(preprocessed_image2)

        predicted_class2 = torch.argmax(outputs2).item()


        probabilities2= torch.nn.functional.softmax(outputs2[0], dim=0)

        # 가장 큰 값의 인덱스
        max_value = max(probabilities2)
        max_index = torch.argmax(probabilities2)

        # 가장 큰 값을 제외한 나머지 값들과 인덱스
        other_values = probabilities2[torch.arange(len(probabilities2)) != max_index]
        other_indices = torch.arange(len(probabilities2))[torch.arange(len(probabilities2)) != max_index]

        # 가장 큰 값과의 차이
        max_difference = max_value - torch.max(other_values)

        # 조건을 충족하는 인덱스 추출
        selected_indices = [max_index.item()]

        # 반복문을 통해 조건을 만족하는 인덱스 찾기
        for idx, value in zip(other_indices, other_values):
            if max_value - value < k:
                selected_indices.append(idx.item())
                
                
        class_names2 = ["A", "B", "C"]
        selected_class_names = [class_names2[idx] for idx in selected_indices]       # k값에 따라 여러개 추출될 수 있는 class
        predicted_class2 = class_names2[predicted_class2]                           # 확률이 가장 높은 값 하나만 추출

        return {'상태': class_names[predicted_class1], "predicted_class": selected_class_names, 'prob': probabilities2.tolist(), '위험도': risk.item()}

    else:

        return {'상태': class_names[predicted_class1], '위험도': risk.item()}

model1을 거친 결과값은 good, normal, bad 세 개의 클래스로 매핑된다. 이때 해당 이미지의 model1 결과가 good이 아닌 경우 즉, normal 혹은 bad인 경우에만 model2로 넘어간다. model2의 task는 해당 이미지가 A, B, C 중에서 어떤 class에 해당하는지 분류하는 것인데, 예를 들어 A class가 가장 높은 확률을 가질 때 그 값이 B, C class에 해당할 확률과 k값 이하로 차이가 나는 것이 있으면 그 class들 또한 출력하도록 하였다.

 

예를 들어 k=0.2로 설정했을 때, model에 softmax를 적용한 결과 A class일 확률이 0.4이고 B class일 확률이 0.3이라면 모델의 출력 결과는 A, B 전부가 된다.

 

 

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)

마지막으로 uvicorn을 실행시켜주면 로컬 서버에 배포가 완료된다.

 

 

 

 

코드를 실행하면 위 결과가 뜨는데, 해당 링크를 클릭한 후, 주소 뒤에 /docs를 붙이면 fastAPI docs 페이지를 확인할 수 있다.

 

 

컴퓨터에 있는 이미지 파일 아무거나 base64 형태로 변환한 후 입력해줬다. base64 특성 상 문자열이 길어도 너무 길어서 용량이 작은 사진이 아니면 입력하는데 굉장히 오래걸린다. base64를 사용할 거라면 파일 크기가 작은 이미지를 사용하는 것이 좋겠다.

 

 

사진을 입력하고 Execute를 누르면 위 같은 결과를 볼 수 있다. 원하는대로 결과가 출력되었으므로 모델 배포 성공~~!!