본문으로 건너뛰기

[ultralytics] COCO Segmentation 검증 300% 속도 향상 — RLE 인코딩 벡터화

PR 링크: ultralytics/ultralytics#22651 상태: Merged | 변경: +101 / -79

들어가며

YOLO 모델의 segmentation 검증(validation) 과정에서 예측된 마스크를 COCO 형식의 JSON으로 변환하려면 Run-Length Encoding(RLE)이 필요합니다. 기존 구현은 faster_coco_eval 라이브러리의 C 확장과 ThreadPool을 사용하여 마스크를 하나씩 인코딩했습니다. 이 방식은 외부 의존성이 필요하고, CPU-GPU 간 데이터 이동이 빈번하며, 스레드 풀 오버헤드까지 발생했습니다.

이 PR은 외부 의존성을 완전히 제거하고, PyTorch tensor 연산으로 배치 RLE 인코딩을 수행하여 검증 속도를 약 300% 향상시킵니다.

핵심 코드 분석

1. ThreadPool + faster_coco_eval 제거

Before:

from faster_coco_eval.core.mask import encode

def single_encode(x):
    """Encode predicted masks as RLE and append results to jdict."""
    rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
    rle["counts"] = rle["counts"].decode("utf-8")
    return rle

pred_masks = np.transpose(predn["masks"], (2, 0, 1))
with ThreadPool(NUM_THREADS) as pool:
    rles = pool.map(single_encode, pred_masks)

After:

pred_masks = predn["masks"].transpose(2, 1).contiguous().view(len(predn["masks"]), -1)  # N, H*W
h, w = predn["masks"].shape[1:3]
counts = multi_encode(pred_masks)
rles = []
for c in counts:
    rles.append({"size": [h, w], "counts": to_string(c)})

기존 코드는 NumPy 배열로 변환 후 C 확장 함수를 스레드 풀로 호출했습니다. 새 코드는 PyTorch tensor를 직접 사용하여 GPU 메모리에서 연산을 수행합니다.

2. multi_encode() — PyTorch 벡터화 RLE

After:

def multi_encode(pixels: torch.Tensor) -> list[int]:
    transitions = pixels[:, 1:] != pixels[:, :-1]
    row_idx, col_idx = torch.where(transitions)
    col_idx = col_idx + 1

    counts = []
    for i in range(pixels.shape[0]):
        positions = col_idx[row_idx == i]
        if len(positions):
            count = torch.diff(positions).tolist()
            count.insert(0, positions[0].item())
            count.append(len(pixels[i]) - positions[-1].item())
        else:
            count = [len(pixels[i])]

        if pixels[i][0].item() == 1:
            count = [0, *count]
        counts.append(count)

    return counts

핵심은 torch.where(transitions)입니다. 인접한 픽셀이 다른 위치(0→1 또는 1→0 전환점)를 한번에 찾아내어, 모든 마스크의 RLE을 배치로 계산합니다. NumPy 루프 대신 PyTorch의 벡터화 연산을 사용하므로 GPU tensor를 그대로 활용할 수 있습니다.

3. scale_image → scale_masks 교체

Before:

"masks": ops.scale_image(
    torch.as_tensor(predn["masks"], dtype=torch.uint8).permute(1, 2, 0).contiguous().cpu().numpy(),
    pbatch["ori_shape"],
    ratio_pad=pbatch["ratio_pad"],
),

After:

"masks": ops.scale_masks(predn["masks"][None], pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"])[
    0
].byte(),

scale_image는 NumPy 배열과 cv2.resize를 사용했습니다. 새 코드는 scale_masks가 F.interpolate(PyTorch)를 사용하므로 GPU tensor를 CPU로 복사할 필요가 없습니다. 또한 scale_masks에 ratio_pad 파라미터를 추가하여 이미 알고 있는 패딩 정보를 재활용합니다.

4. to_string() — Delta + Variable-Length Encoding

def to_string(counts: list[int]) -> str:
    result = []
    for i in range(len(counts)):
        x = int(counts[i])
        if i > 2:
            x -= int(counts[i - 2])
        while True:
            c = x & 0x1F  # Take 5 bits
            x >>= 5
            more = (x != -1) if (c & 0x10) else (x != 0)
            if more:
                c |= 0x20  # Set continuation bit
            c += 48  # Shift to ASCII
            result.append(chr(c))
            if not more:
                break
    return "".join(result)

COCO RLE 문자열 형식을 순수 Python으로 구현합니다. Delta encoding(3번째 이후 카운트는 2칸 앞 값과의 차이로 저장)과 variable-length encoding(5비트씩 나누어 ASCII 문자로 변환)을 조합하여 컴팩트한 문자열을 생성합니다.

왜 이게 좋은가

  1. 외부 의존성 제거: faster_coco_eval 패키지가 더 이상 필요하지 않아 설치 과정이 단순해졌습니다.
  2. GPU→CPU 전환 최소화: 마스크 데이터가 PyTorch tensor 형태로 GPU에 머물면서 처리되어 불필요한 메모리 복사를 줄입니다.
  3. 스레드 풀 오버헤드 제거: ThreadPool 생성/관리 비용과 GIL 경합이 사라집니다.
  4. 벡터화 연산: 개별 마스크 루프 대신 torch.where로 모든 전환점을 한번에 찾습니다.

정리

  • GPU tensor를 최대한 오래 유지하라: CPU/GPU 간 데이터 이동은 숨겨진 성능 병목입니다. torch.as_tensor(...).cpu().numpy() 같은 체인은 경고 신호입니다.
  • 외부 C 확장 vs 벡터화 Python: 충분히 벡터화된 PyTorch 코드는 개별 호출되는 C 확장보다 빠를 수 있습니다. 배치 크기가 클수록 차이가 커집니다.
  • ThreadPool은 만능이 아니다: GIL이 존재하는 Python에서 CPU-bound 작업의 ThreadPool은 실질적 병렬성을 제공하지 못합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글