[Pytorch] 모델 서빙하기

[원본 링크]

만들어진 모델 파일을 가져다가 사용해보는 방법을 다뤄본다.
학습 과정 같은 것은 다루지 않는다.

여기서 사용해볼 모델은 resnet-18이라고 하는, 이미지 분류 모델이다.
https://huggingface.co/microsoft/resnet-18
이미지를 받으면 그게 어떤 형태의 이미지인지 분석해서 텍스트 목록으로 알려주는 기능을 갖고 있다.

예시로 사용한 테스트코드는 깃헙을 참조
https://github.com/myyrakle/ml_examples/tree/master/pytorch/start-with-resnet-18




종속성 설치 (uv)

보통 이 정도는 기본으로 깔고 간다고 생각하면 된다.
pillow는 이미지 처리 모듈, torch가 pytorch 모듈, transformer가 모델을 로드해서 쏘기 위한 모듈이다.

uv add pillow torch transformers




기본 세팅

먼저 모델부터 로딩해보자.

from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image

# 모델 로딩
image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-18")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18")

# ID과 레이블 값 매핑 조회
print(model.config.id2label)

기본적으로 이런 모델들은, 출력의 형태가 정해져있다.

그래서 모델 자체에 class id와 label name의 매핑이 존재하고, 모델 추론의 결과로는 class id가 반환되어서 그걸로 이 config에서 값을 찾아서 사용하는 식이 되는 것이다.


이제 추론을 돌려보자.
먼저 이미지를 하나 준비하고

이번에는 추론 코드까지 작성해보겠다.

from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image

# 모델 로딩
image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-18")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18")

# 이미지 넣어서 입력 생성
image = Image.open("cat.jpg")
inputs = image_processor(image, return_tensors="pt")

print(model.config.id2label)

# gradient 계산 비활성화한 채로 모델 추론 실행
with torch.no_grad(): 
    logits = model(**inputs).logits

# 모델 결과값 출력
predicted_label = logits.argmax(-1).item() # 가장 높은 점수를 가진 클래스의 인덱스
print(f"Predicted label: {predicted_label}")
print(model.config.id2label[predicted_label])

logits의 결과값으로 추론된 class 인덱스들이 반환되고, 그걸로 id 매핑을 뒤져서 값을 가져왔다.


그러면 tabby, tabby cat이라는 값이 추론된 것을 볼 수 있다.

혹은, 상위 N개를 가져와서 처리하는 식으로 할 수도 있다.
애초에 이런 추론 결과는 정확하지 않고 그냥 점수가 높은 것을 가져오는 식으로 동작하기 때문이다.

# 상위 k개 클래스 보기
k = 5  # 상위 5개
top_k = torch.topk(logits, k)
print(f"\n상위 {k}개 예측:")
for i in range(k):
    idx = top_k.indices[0][i].item()
    score = top_k.values[0][i].item()
    label = model.config.id2label[idx]
    print(f"{i+1}. {label}: {score:.4f}")

그러면 이런 식으로 응답이 나온다. score가 추론된 점수다.




모델 파일

위에서는 모델 파일을 명시적으로 다운받지 않고, 그냥 레포지토리 경로로 로드했었다.

image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-18")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18")

이런 식으로 치면 transformers 모듈은 huggingface 레포지토리에서 모델을 찾아서 자동으로 다운받는다.
처음 시작하기에는 간편하지만, 모델을 huggingface에 올리지 않을 수도 있고, 아무튼 로컬에 모델을 직접 다운받아서 쓰고 싶을 수도 있다.

모델 자체는 huggingface cli를 통해 설치할 수 있다.

pip install huggingface-hub
huggingface-cli download microsoft/resnet-18 --local-dir ./resnet-18

그러면 디렉터리 기반으로 이렇게 이것저것 다운이 된다.

안에 뭐가 들어있는지는 모델마다 다르다.
하지만 대체로 config.json이 모델 구조를 정의하는 파일, preprocessor_config.json이 전처리 설정파일.

그리고 이 경우에는 model.safetensors가 프레임워크 독립적인 모델 파일이다. 그냥 기본 옵션으로 로드할 경우에는 model.safetensors만 있어도 된다.
참고로 tf_model.h5가 텐서플로 모델이고, pytorch_model.bin은 파이토치 모델(.pt)이다.

그러면 이제 저 디렉터리를 기반으로 로컬에 존재하는 모델을 로드해서 쓸 수 있고

마찬가지로 잘 동작할 것이다.

use_safetensor=False 옵션을 넣으면 강제로 pytorch_model.bin을 사용하도록 할 수 있다.

저 pytorch_model.bin이 일반적으로 말하는 .pt 파일이다.

safetensor를 만든 huggingface 측에서는 일반적인 모델 pickle 형식보다 safetensor를 쓰기를 권장한다. 훨씬 보안에 안전하고 로딩 등이 빠르다는 이유에서다.
이 포맷들에 대해서는 별도 포스트에서 따로 다뤄보겠다.



참조
https://wikidocs.net/201325