TorchServe

TorchServe

TorchServe는 PyTorch 모델을 프로덕션에 배포하고 서빙하기 위한 도구예요. 학습된 PyTorch 모델을 압축 포장(MAR)해 저장하고, 추론 API·관리 API를 통해 서빙할 수 있게 해줘요. REST와 gRPC API를 지원해서, 이미지 분류부터 시퀀스 생성까지 다양한 모델을 일관된 방식으로 내보내요.

출처: PyTorch/Serve 문서 (공식)

시작하기

TorchServe와 모델 아카이버를 설치하는 것부터 시작해요. Conda와 pip 모두 지원돼요.

conda install torchserve torch-model-archiver torch-workflow-archiver -c pytorch

pip으로는 이렇게 설치해요.

pip install torchserve torch-model-archiver torch-workflow-archiver

TorchServe에서 모델을 서빙하려면 먼저 모델을 MAR 파일로 아카이빙해야 해요. 모델 아카이버로 모델을 패키징하면 됩니다.

모델 저장하기

모델을 저장할 디렉토리를 만들고, 학습된 모델 파일을 받아서 아카이버로 포장해요. 이 예제는 DenseNet161 이미지 분류기를 MAR로 만드는 흐름이에요.

mkdir model_store
wget https://download.pytorch.org/models/densenet161-8d451a50.pth

모델 파일, 핸들러(image_classifier), 클래스 이름 매핑 파일(index_to_name.json)을 함께 묶어요.

torch-model-archiver --model-name densenet161 --version 1.0 \
  --model-file ./serve/examples/image_classifier/densenet_161/model.py \
  --serialized-file densenet161-8d451a50.pth \
  --export-path model_store \
  --extra-files ./serve/examples/image_classifier/index_to_name.json \
  --handler image_classifier

서버 시작하기

MAR 파일을 저장한 뒤 torchserve 명령으로 모델을 서빙해요.

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

기본 추론 핸들러

TorchServe는 상자에서 꺼내 바로 쓸 수 있는 기본 추론 핸들러들을 제공해요.

  • image_classifier — ImageNet 데이터로 학습한 이미지 분류 모델 처리. RGB 이미지를 입력받아 top-5 예측과 확률을 출력해요.
  • image_segmenter — 이미지 세그멘테이션 모델 처리. [N, CL, H, W] 형태의 출력을 내요.
  • object_detector — 객체 탐지 모델 처리. 탐지된 클래스와 바운딩 박스 목록을 출력해요.
  • text_classifier — AG_NEWS 데이터로 학습한 텍스트 분류 모델 처리. (배칭 미지원)

index_to_name.json에 클래스 번호(문자열)에서 이름으로의 매핑을 담으면, image_classifier·text_classifier·object_detector 핸들러가 숫자 클래스를 친근한 문자열로 자동 매핑해 줘요.

더 알아보기