본문 바로가기

공부방/Upstage AI Lab 4기

[CV 문서 분류 대회] 배치 크기(batch size)는 어떻게 결정해야 하나

배치 사이즈는 실험하는 사람이 결정할 수 있는 파라미터 중 하나로, 최적값은 아래 3가지 요소를 보고 결정해야 한다. 

1) GPU 메모리 크기: 가장 중요한 제한 요소
2) 모델 크기: 더 큰 모델은 더 많은 메모리 필요
3) 이미지 크기: 큰 이미지는 더 많은 메모리 사용

GPU 메모리 크기를 넘으면 RuntimeError: CUDA out of memory 이런 에러가 뜬다.

일반적으로 배치 크기는 16, 32, 64, 128, 256 이런 식으로 많이 쓰는데, 실험적으로 접근해보면 좋다.

# 배치 크기별 실험 예시
batch_sizes = [16, 32, 64, 128]
results = {}

for bs in batch_sizes:
    BATCH_SIZE = bs
    trn_loader = DataLoader(
        trn_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True
    )
    # 모델 학습 및 결과 기록
    results[bs] = train_model()

 

트레인 데이터셋에는 1570개의 이미지가 있고, 기본 베이스라인에서 데이터셋의 이미지 크기를 32*32로 변환해 놓은 상태이다.

에폭이 한 번 돌 때마다 생성되는 배치의 개수는 (트레인 데이터셋 개수) / (배치 사이즈)이므로,
배치를 32라고 뒀을 때 1570/ 32 = 약 49개 배치가 생성된다. 

트레인 데이터셋을 크게 늘린다면, 배치 사이즈를 키워도 괜찮다. 

이미지 사이즈를 크게 증가시키면 메모리의 사용량이 크게 증가하기 때문에 배치 사이즈를 줄여야 할 수도 있다. (확인해보지는 않음) 

# 이미지 한 장당 메모리 사용량 (채널=3 기준)
# 32x32 이미지
memory_32 = 32 * 32 * 3 = 3,072 픽셀

# 260x260 이미지
memory_260 = 260 * 260 * 3 = 202,800 픽셀

# 메모리 증가율
increase = memory_260 / memory_32 = 66배

 

배치 사이즈가 작아지면, 

  • 한 에포크당 더 많은 업데이트 발생
  • 학습이 더 오래 걸릴 수 있음
  • 일반화 성능이 좋아질 수 있음

배치 사이즈에 따라 학습률을 조정해야 할 수 있다. 

음? 배치 사이즈랑 학습률은 무슨 연관성이 있길래...? 

# 일반적인 경험적 법칙
# 배치 크기가 2배가 되면 학습률도 비례해서 증가
BATCH_SIZE = 32
LR = 1e-3

# 배치 크기를 2배로 증가시킬 경우
BATCH_SIZE = 64
LR = 2e-3  # 학습률도 2배로 증가

배치 사이즈가 더 크면 더 안정적인 그래디언트를 갖기 때문에 더 큰 학습률을 사용할 수 있고,
반대로 배치가 작으면 그래디언트에 노이즈가 많아서 더 작은 학습률이 필요하다고.

 

에포크는 배치 사이즈가 작으면, 하나의 에포크 안에서 업데이트가 많이되는데
배치 사이즈가 크면 한 에포크당 업데이트가 적어서 에포크 양을 늘리는 게 나을 수도 있다. 
1570개 이미지 가지고는 배치 사이즈 32일 때 한 에폭당 약 49회(1570/32)를 업데이트하고,
1570개 이미지 가지고 배치 사이즈 64일 때 한 에폭당 약 25회(1570/64)를 업데이트함.