이 글은 알고리즘 팀에서 작성한 모델 파일을 통해 이미지를 추론하여 그 결과를 가공하고 DB에 적재하는 업무를 하던 중 경험한 내용을 정리한 글입니다.
먼저 알고리즘 팀으로부터 모델 파일과 추론 결과값 형식을 제공받은 후 서빙 코드를 작성했습니다.
def inference_on_gpu(device_id: int, image_list: List[str]) -> None:
try:
# 모델 로드
model = load_model(model_path=model_path, device_id=device_id, save_dir=save_dir, save_folder="run")
for jpg_filename in image_list:
# 이미지 하나씩 모델 실행
results = model.predict(
jpg_filename,
save=True, # 추론 결과를 저장할지 여부 O
)
# 기타 로직 수행...
# txt 파일에 결과 쓰기
txt_filename = change_extension(jpg_filename, "txt")
with open(txt_filename, "w") as f:
for result in results:
f.write(...)
except Exception as e:
# 예외 처리
def main():
# 기타 로직 수행...
# 가용 가능한 GPU 수 파악
gpu_list = [int(d.strip()) for d in args.gpus.split(",") if d.strip() != ""]
# jpg 파일 검색
jpg_paths = find_files_by_extension(directory_path=...)
# jpg 파일을 GPU 수에 맞게 균등 분할
partitions = partition_list(jpg_paths, len(gpu_list))
# 가용 가능한 GPU 병렬 처리
with ProcessPoolExecutor(max_workers=len(gpu_list)) as executor:
futures = []
for device, partition in zip(gpu_list, partitions):
futures.append(
executor.submit(inference_on_gpu, device, partition, ...)
)이 코드의 프로세스에 대하여 간략하게 설명하면 다음과 같습니다.
- 가용 GPU 수 확인
- 입력 이미지 수집
- 이미지를 GPU 수에 따라 균등하게 분할
- GPU 병렬 처리로 YOLO 모델 추론
- 추론 결과를 텍스트 파일로 저장
위 프로세스 중 1번부터 4번까지는 개선할 수 있는 포인트가 없다고 판단했습니다.
하지만 5번의 경우 여러 장의 이미지를 하나하나 순회하면서 모델을 추론하고 텍스트 파일에 결과를 저장하고 있어서, 개선 가능성을 찾기 시작했습니다.
그 결과 다음과 같은 사실을 알게 되었습니다.
- GPU 유휴 시간 발생 : 추론이 끝난 후 CPU I/O 작업 완료까지 다음 추론이 대기
- GPU-CPU 메모리 전송 비용 : GPU 메모리에서 CPU 메모리로의 데이터 복사 오버헤드 발생
저는 GPU 유휴 시간을 최소화하기 위해 추론 결과 파일 쓰기를 비동기 처리해야겠다는 생각이 들었고, 그 과정에서 큐 자료구조와 스레드를 사용하여 더 효과적으로 성능을 개선했습니다.
구조는 다음과 같습니다.
먼저 GPU에서 추론이 완료된 결과를 Queue에 보내고 바로 다음 이미지 추론을 시작하고, 스레드는 Queue에서 결과값을 꺼내어 쓰기 작업을 하는 방식입니다.
Queue 관련 코드
write_queue = Queue(maxsize=200) # 큐 생성
for i, image_path in enumerate(image_list):
# 이미지 추론
write_queue.put((txt_filename, txt_lines), timeout=0.1) # 추론 결과 Queue에 삽입
Thread 관련 코드
""" Thread에서 실행될 파일 쓰기 로직 """
def file_writer_thread(write_queue: Queue, device_id: int):
written_count = 0
while True:
try:
# 큐에서 작업 가져오기 (타임아웃 1초)
item = write_queue.get(timeout=1)
# 종료 신호 확인
if item is None:
break
filename, lines = item
# 파일 쓰기 실행
try:
with open(filename, "w") as f:
f.writelines(lines)
written_count += 1
# 예외 처리 ...
# Thread에서 실행할 함수 지정
writer_thread = Thread(target=file_writer_thread, args=(write_queue, device_id))
writer_thread.daemon = True # 백그라운드 작업 진행
writer_thread.start() # Thread 시작 포인트
결과
2,048 x 12,000 사이즈의 이미지 12,780개를 처리하는데 1426.81초가 소요되던 원래 코드를 수정하여 672.75초로 단축시켜 약 `52.85%` 성능 개선을 달성했습니다.