본격적으로 Segment Anything 모델을 코드와 함께 Torchserve 배포하는 과정을 살펴보도록 하겠습니다. 이번 글은 코드 중심으로 진행될 예정이니 Torchserve에 대해 좀 더 자세한 설명이 궁금하시다면 저희 블로그의 Torchserve 시리즈를 읽어보시면 좋을 것 같습니다. 이제 시작해보도록 하겠습니다.
Torchserve handler 작성하기
추론을 위해 필요한 config.properties와 handler.py를 작성해보겠습니다.
config.properties의 address port는 이후 docker 실행 시 포트 포워딩을 해 docker 밖에서도 추론할 수 있도록 설정해 줄 예정입니다. gpu_id와 batch_size는 handler.py의 properties:dict 변수로 들어가는 값입니다.
handler.py 기본적인 코드 구조는 Torchserve에서 제공하는 예시 코드와 동일합니다. 대신 후처리에서 추론 결과를 base64로 인코딩하는 과정을 추가해주었습니다. torch.Tensor나 np.array 그대로 반환해도 상관은 없지만, 이후 편리한 통신을 위해 string 타입으로 내보내겠습니다.
먼저 DockerFile을 작성해보겠습니다. Base image는 원하시는 버전을 선택하시면 되고, Torchserve 사용 시 jdk가 필요하기 때문에 설치해주겠습니다. 마지막으로 필요한 라이브러리를 설치해주면 끝입니다. 편이에 따라 requirements.txt 파일을 만들어 RUN pip install -r requirements.txt로 설치하셔도 됩니다.
FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
WORKDIR /workspace
RUN mkdir datahunt_segment_anything
# Install OpenJDK-11
RUN apt-get update && \
apt-get install -y git && \
apt-get install -y openjdk-11-jre-headless && \
apt-get clean;
RUN apt-get -y install libgl1-mesa-glx && \
apt-get install -y curl;
ADD datahunt_segment_anything /workspace/datahunt_segment_anything
RUN pip install git+https://github.com/facebookresearch/segment-anything.git
RUN pip install opencv-python matplotlib onnxruntime onnx
RUN pip install torchserve torch-model-archiver torch-workflow-archiver nvgpu validators tensorflow-cpu
이제 빌드해보도록 하겠습니다.
docker build -t segment-anything:v1 .
다음으로 컨테이너를 만들고 포트 설정을 해주겠습니다. 필요에 따라 volume 마운팅 설정을 해주면 로컬에서 작업하는 내용을 실시간으로 반영할 수 있습니다. 하지만 저희는 완성된 코드를 사용할 예정이니 따로 설정하진 않겠습니다. -p 8070은 위에 config.properties에서 설정해준 inference_address의 port 입니다.
docker run -it --gpus all --name segment-anything-v1 -p 8070:8070/tcp segment-anything:v1 /bin/bash
컨테이너에 접속한 걸 확인하셨나요? 그럼 이제 배포를 위한 모든 준비가 끝났으니 .mar 파일을 생성하고 Torchserve를 실행보겠습니다.
Segment Anything 모델 배포하기
현재 작업 디렉토리는 /workspace/datahunt_segment_anything이고, 내부 폴더 구조는 다음과 같습니다.
배포를 위한 스크립트는 아래와 같습니다. 직접 하나씩 실행해도 괜찮지만 이후 반복 사용을 위해 쉘 스크립트로 작성하겠습니다. 재배포 용도로 사용할 예정이기 때문에 실행 중인 서버가 있다면 종료하고, 기존에 만들어진 mar도 삭제하는 코드를 추가했습니다. 해당 과정이 필요하지 않다면 삭제해도 괜찮습니다.
#!/bin/bash
torchserve --stop
version=1.0
model_name="SAM"
model_path='../model/sam_vit_l.pth'
if [ -e ${model_name}.mar ]; then
rm ${model_name}.mar
echo 'Removed existing model archive.'
fi
# Create mar file
torch-model-archiver --model-name ${model_name} --version ${version} --serialized-file ${model_path} --handler "handler.py"
# Depolyment
torchserve --start --model-store . --models ${model_name}.mar --ts-config ./config.properties
이제 위에서 작성한 torchserve_deploy.sh 파일을 실행하면 모델 배포가 완료됩니다. 모델 추론을 해볼까요? 원하는 샘플을 준비하고 아래와 같이 실행해 보겠습니다.
curl -X POST "http://127.0.0.1:8070/predictions/SAM/1.0" -T "test/sample/sample.jpg"
그럼 아래와 같이 handler.py에서 설정한 형식대로 결과가 출력된 것을 확인할 수 있습니다.
handler 추론 결과 (base64 encoding)
Segment Anything Promptable Task 추론 테스트
미리 뽑은 임베딩 결과를 불러와 Promptable Task를 진행하는 과정입니다. Point와 Bounding box 좌표는 미리 가지고 있던 값을 사용했습니다. 여기서는 Python 코드를 사용했지만 Github Demo에 다른 언어를 사용해 추론하는 코드도 있으니 참고해주세요.
.svg)
