[ultralytics] Ultralytics 체크포인트 로딩 최적화 및 스레드 안전성 강화

PR 링크: ultralytics/ultralytics#25885 상태: Merged | 변경: +57 / -32

들어가며

최근 Ultralytics 레포지토리의 체크포인트 로딩 메커니즘에서 중요한 개선이 이루어졌습니다. 기존의 weights_only=True를 활용한 restricted loading 방식은 멀티스레드 환경에서 allow-list를 관리할 때 스레드 간 충돌이 발생할 수 있는 잠재적 위험이 있었고, 불필요하게 큰 allow-list를 매번 처리함으로써 로딩 속도가 저하되는 문제가 있었습니다. 본 PR은 이러한 문제를 해결하기 위해 프로세스 전역 락(process-wide lock)을 도입하고, 체크포인트별로 필요한 글로벌 항목만 동적으로 등록하도록 최적화했습니다.

코드 분석

1. ultralytics/nn/tasks.py: 스레드 안전한 체크포인트 로딩

기존에는 safe_globals 컨텍스트 매니저가 프로세스 전역 상태를 변경하며 스레드 간 레이스 컨디션을 유발했습니다. 이를 해결하기 위해 threading.Lock을 도입하고, 필요한 항목만 등록하도록 변경했습니다.

Before:

# 기존 방식: 전체 리스트를 매번 등록하고 스레드 간 간섭 발생 가능
with torch.serialization.safe_globals(allow):
    yield

After:

# 개선된 방식: 필요한 항목만 식별하여 락을 통해 안전하게 등록
needed = torch.serialization.get_unsafe_globals_in_checkpoint(weight)
with cls._lock:
    # ... registry 구축 및 필요한 항목만 add_safe_globals 호출
    if entries:
        torch.serialization.add_safe_globals(entries)

2. ultralytics/nn/modules/block.py: RLE Prior 계산 최적화

MultivariateNormal 객체를 매번 생성하는 대신, 2차원 표준 정규 분포의 로그 확률을 직접 계산하는 방식으로 변경하여 오버헤드를 줄였습니다.

Before:

@property
def prior(self):
    return torch.distributions.MultivariateNormal(self.loc, self.cov)

After:

# Closed-form log N(z; 0, I) 계산으로 대체
return -0.5 * (z.float() ** 2).sum(-1) - math.log(2 * math.pi) + log_det

왜 이게 좋은가

1. 성능 향상: 체크포인트가 참조하는 글로벌 항목만 등록함으로써, 660개에 달하는 전체 allow-list를 매번 처리할 필요가 없어졌습니다. 이로 인해 로딩 속도가 약 40% 향상되었습니다.
2. 스레드 안전성: threading.Lock을 사용하여 멀티스레드 환경에서 여러 체크포인트가 동시에 로드될 때 allow-list가 서로 덮어쓰여지는 문제를 방지했습니다. test_restricted_load_threaded 테스트를 통해 32개의 로드 작업을 8개의 스레드에서 동시에 수행해도 안전함을 검증했습니다.
3. 지연 로딩(Lazy Loading): torchvision과 같은 무거운 라이브러리는 체크포인트가 실제로 해당 변환(transform)을 필요로 할 때만 임포트되도록 하여, 불필요한 의존성 로드를 제거했습니다.

이러한 최적화는 대규모 모델 배포 환경에서 체크포인트 로딩의 안정성과 속도를 동시에 확보해야 할 때 매우 중요한 교훈을 줍니다. 특히 torch.serialization과 같은 저수준 API를 다룰 때는 전역 상태 변경에 따른 부작용을 항상 고려해야 합니다.

참고 자료

• torch.serialization.add_safe_globals
• torch.serialization.get_unsafe_globals_in_checkpoint

참고 자료

• https://pytorch.org/docs/stable/generated/torch.serialization.add_safe_globals.html
• https://pytorch.org/docs/stable/generated/torch.serialization.get_unsafe_globals_in_checkpoint.html

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