Skip to content

MMSegmentation

Triton Inference Server 로 올리기

MMSegmentation 모델을 NVIDIA Triton Inference Server에서 사용하는 방법을 단계별로 설명드리겠습니다. MMSegmentation은 PyTorch를 기반으로 하는 모델이며, Triton Inference Server와 함께 사용하려면 ONNX 또는 TorchScript 형식으로 모델을 변환해야 합니다. 이후 Triton 서버에 모델을 배포하고 추론을 실행할 수 있습니다.

ONNX 로 모델 변환

먼저 MMSegmentation 모델을 ONNX 또는 TorchScript로 변환합니다. 

import torch
from mmseg.apis import init_segmentor
import mmcv

# 모델과 구성 파일 경로
config_file = 'configs/fcn/fcn_r50-d8_512x512_40k_voc12aug.py'
checkpoint_file = 'checkpoints/fcn_r50-d8_512x512_40k_voc12aug_20200617_071005-41a67fa4.pth'

# 모델 초기화
model = init_segmentor(config_file, checkpoint_file, device='cuda:0')

# 입력 데이터 생성 (예시 입력)
dummy_input = torch.randn(1, 3, 512, 512).cuda()

# 모델을 평가 모드로 전환
model.eval()

# 모델 변환
torch.onnx.export(model, dummy_input, "mmsegmentation_model.onnx", 
                  export_params=True, opset_version=11, 
                  input_names=['input'], output_names=['output'])

모델 리포지토리 구성

모델을 Triton Inference Server의 모델 리포지토리에 배치합니다. 디렉토리 구조는 다음과 같습니다: 

<model-repository>
└── mmsegmentation_model
    ├── 1
    │   └── model.onnx
    └── config.pbtxt

config.pbtxt 수정: 

name: "mmsegmentation_model"
platform: "onnxruntime_onnx"
max_batch_size: 1
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [ 3, 512, 512 ]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [ 512, 512 ]
  }
]

클라이언트에서 추론 요청

이후 tritonclient 를 사용하여 Triton Inference Server에 추론 요청을 보냅니다. 필요에 따라 공유 메모리를 사용하여 성능을 최적화할 수 있습니다. 

import tritonclient.http as httpclient
import tritonclient.utils.shared_memory as shm
import numpy as np

# Triton Inference Server URL
url = "localhost:8000"

# 클라이언트 초기화
client = httpclient.InferenceServerClient(url=url)

# 모델 정보
model_name = "mmsegmentation_model"
input_name = "input"
output_name = "output"

# 입력 데이터 준비
input_data = np.random.rand(1, 3, 512, 512).astype(np.float32)

# Shared memory 설정
input_byte_size = input_data.nbytes
output_byte_size = input_byte_size  # 예제에서는 입력과 출력 크기가 같다고 가정

# Shared memory 생성
input_shm_handle = shm.create_shared_memory_region("input_data", "/dev/shm/input", input_byte_size)
output_shm_handle = shm.create_shared_memory_region("output_data", "/dev/shm/output", output_byte_size)

# Shared memory에 데이터 쓰기
shm.set_shared_memory_region(input_shm_handle, [input_data])

# Shared memory 등록
client.register_system_shared_memory("input_data", "/dev/shm/input", input_byte_size)
client.register_system_shared_memory("output_data", "/dev/shm/output", output_byte_size)

# 입력 및 출력 설정
inputs = []
outputs = []

inputs.append(httpclient.InferInput(input_name, input_data.shape, "FP32"))
inputs[-1].set_shared_memory("input_data", input_byte_size)

outputs.append(httpclient.InferRequestedOutput(output_name, binary_data=True))
outputs[-1].set_shared_memory("output_data", output_byte_size)

# 추론 요청
results = client.infer(model_name, inputs, outputs=outputs)

# 결과 읽기
output_data = shm.get_shared_memory_region(output_shm_handle, output_byte_size, np.float32)
print(output_data)

# Shared memory 해제
client.unregister_system_shared_memory("input_data")
client.unregister_system_shared_memory("output_data")
shm.destroy_shared_memory_region(input_shm_handle)
shm.destroy_shared_memory_region(output_shm_handle)

See also