[Ultralytics] v8.4.28 — autobatch를 데이터셋 크기로 제한하여 소규모 데이터셋 학습 안정화

PR 링크: ultralytics/ultralytics#24020 상태: Merged | 변경: +48 / -45

들어가며

GPU 메모리가 충분한 환경에서 작은 데이터셋을 학습할 때, autobatch가 데이터셋 전체 크기보다 큰 배치를 산출하는 비효율이 발생할 수 있습니다. 예를 들어 50장의 이미지로 학습할 때 autobatch가 64를 반환하면, DataLoader가 매 에포크마다 불필요한 빈 배치를 생성하거나 drop_last로 인해 데이터를 소모합니다. 이번 PR은 autobatch 결과를 데이터셋 크기로 상한을 설정하고, INT8 캘리브레이션의 배치 크기 처리도 개선합니다.

핵심 코드 분석

1. autobatch 함수에 dataset_size 파라미터 추가

Before:

def autobatch(model, imgsz=640, fraction=0.60, batch_size=16, max_num_obj=1):
    batch_sizes = [1, 2, 4, 8, 16] if t < 16 else [1, 2, 4, 8, 16, 32, 64]

After:

def autobatch(model, imgsz=640, fraction=0.60, batch_size=16, max_num_obj=1, dataset_size=0):
    batch_sizes = [1, 2, 4, 8, 16] if t < 16 else [1, 2, 4, 8, 16, 32, 64]
    if dataset_size > 0:
        batch_sizes = [b for b in batch_sizes if b <= dataset_size]
    # ... 프로파일링 후
    if dataset_size > 0:
        b = min(b, dataset_size)

두 곳에서 dataset_size를 적용합니다. 먼저 프로파일링 대상 배치 크기 리스트를 데이터셋 크기 이하로 필터링하고, 최종 산출 값도 상한을 적용합니다. 이로써 불필요하게 큰 배치로 메모리를 프로파일링하는 낭비도 방지됩니다.

2. detect/train에서 dataset_size 전달

Before:

def auto_batch(self):
    train_dataset = self.build_dataset(self.data["train"], mode="train", batch=16)
    max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4
    del train_dataset
    return super().auto_batch(max_num_obj)

After:

def auto_batch(self):
    train_dataset = self.build_dataset(self.data["train"], mode="train", batch=16)
    max_num_obj = max(len(label["cls"]) for label in train_dataset.labels) * 4
    n = len(train_dataset)
    del train_dataset
    return super().auto_batch(max_num_obj, dataset_size=n)

데이터셋을 빌드하는 시점에서 이미 len(train_dataset)을 알 수 있으므로, 추가 비용 없이 dataset_size를 전달합니다.

3. INT8 캘리브레이션 배치 크기 유연화

Before:

if n < self.args.batch:
    raise ValueError(
        f"The calibration dataset ({n} images) must have at least as many images "
        f"as the batch size ('batch={self.args.batch}')."
    )

After:

if n < 1:
    raise ValueError(f"The calibration dataset must have at least 1 image, but found {n} images.")
batch = min(self.args.batch, n)
if n < self.args.batch:
    LOGGER.warning(
        f"{prefix} calibration dataset has only {n} images, reducing calibration batch size to {batch}."
    )
return build_dataloader(dataset, batch=batch, workers=0, drop_last=True)

기존에는 캘리브레이션 데이터셋이 배치 크기보다 작으면 에러를 발생시켰습니다. 이제는 자동으로 배치 크기를 줄이고 경고만 출력합니다. 이로써 적은 수의 캘리브레이션 이미지로도 INT8 export가 가능해집니다.

왜 이게 좋은가

1. 소규모 데이터셋 친화적: PoC나 파인튜닝에서 흔한 소규모 데이터셋(수십~수백 장)에서의 학습 경험이 크게 개선됩니다.

2. 에러에서 경고로 전환: INT8 캘리브레이션에서 hard error 대신 자동 조정 + 경고로 전환하여, 사용자가 수동으로 배치 크기를 조정할 필요가 없어졌습니다.

3. 프로파일링 효율성: 데이터셋보다 큰 배치 크기를 프로파일링하지 않아 autobatch 탐색 시간도 단축됩니다.

정리

dataset_size 파라미터를 autobatch 체인에 전파하여 배치 크기 상한을 설정하고, INT8 캘리브레이션의 배치 크기 제약을 완화한 릴리스(v8.4.28) PR입니다. 소규모 데이터셋에서의 학습과 양자화 워크플로우를 모두 개선합니다.

참고 자료

• ultralytics/ultralytics#24020
• Ultralytics autobatch 문서

---

이 글은 AI(Claude)의 도움을 받아 작성되었으며, 실제 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.