본문으로 건너뛰기

[ultralytics] Ultralytics 추론 성능 극대화: CPU 전처리 병목 해소하기

PR 링크: ultralytics/ultralytics#25982 상태: Merged | 변경: +35 / -14

들어가며

딥러닝 모델의 추론 파이프라인에서 데이터 전처리는 종종 간과되기 쉬운 병목 구간입니다. 특히 CPU에서 NumPy를 사용하여 이미지를 변환(BGR to RGB, NHWC to NCHW)하고 np.ascontiguousarray()를 호출하는 과정은 GPU의 빠른 연산 속도를 따라가지 못해 전체 파이프라인의 처리량을 제한합니다. 이번 Ultralytics PR은 이러한 CPU 기반의 전처리 로직을 디바이스(GPU) 텐서 연산으로 이동시켜, 데이터 전송 및 변환 오버헤드를 획기적으로 줄였습니다.

코드 분석

1. ultralytics/engine/predictor.py의 전처리 로직 개선

기존 방식은 CPU에서 모든 전처리를 마친 후 텐서로 변환하여 GPU로 전송했습니다. 하지만 이번 변경을 통해 uint8 상태의 데이터를 빠르게 GPU로 전송한 뒤, GPU의 병렬 연산 능력을 활용해 차원을 변경하고 정규화를 수행합니다.

Before:

im = np.stack(self.pre_transform(im))
if im.shape[-1] == 3:
    im = im[..., ::-1]  # BGR to RGB
im = im.transpose((0, 3, 1, 2))
im = np.ascontiguousarray(im)
im = torch.from_numpy(im).to(self.device)

After:

im = torch.from_numpy(np.stack(self.pre_transform(im)))
im = im.to(self.device)  # transfer as uint8
im = im.permute(0, 3, 1, 2)  # BHWC to BCHW
if im.shape[1] == 3:
    im = im.flip(1)  # BGR to RGB
im = im.contiguous()
im = (im.half() if self.model.fp16 else im.float()).div_(255)

핵심은 np.ascontiguousarray와 같은 CPU 연산을 제거하고, permute, flip, div_와 같은 PyTorch의 In-place 연산을 활용하여 메모리 복사를 최소화한 점입니다.

2. 회귀 테스트 추가 (tests/test_python.py)

최적화 과정에서 수치적 정밀도가 깨지지 않도록 bit-exact 테스트를 도입했습니다. NumPy 참조 파이프라인과 텐서 연산 결과가 동일한지 검증하여 최적화의 안정성을 확보했습니다.

왜 이게 좋은가

성능 향상

RTX 5080 환경에서 벤치마크 결과, 배치 사이즈 32 기준으로 약 3.11배의 속도 향상을 보였습니다. 이는 CPU에서 수행하던 무거운 배열 조작을 GPU의 고속 메모리 대역폭을 활용하여 처리했기 때문입니다.

기술적 교훈

  1. 데이터 전송 최소화: CPU와 GPU 간의 데이터 이동은 비용이 큽니다. 가능한 한 '가벼운' 상태(uint8)로 데이터를 넘긴 뒤, GPU에서 변환하는 것이 유리합니다.
  2. In-place 연산 활용: .div_()와 같은 In-place 연산은 새로운 메모리 할당을 방지하여 메모리 효율을 높입니다.
  3. 디바이스 특수성 고려: 리뷰 과정에서 논의된 것처럼, non_blocking=True 전송은 MPS(Apple Silicon)와 같은 특정 하드웨어에서 데이터 오염을 유발할 수 있습니다. 최적화 시 타겟 디바이스의 아키텍처 특성을 고려하는 것이 필수적입니다.

결론

이번 개선은 단순히 코드를 옮기는 것을 넘어, 연산의 주체를 CPU에서 GPU로 전환함으로써 현대적인 하드웨어 가속기를 최대한 활용하는 모범 사례를 보여줍니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글