[Machine Learning] 예시를 통해 살펴보는 배치 사이즈(Batch size)의 이해

머신러닝 모델을 학습시키는 과정엔 다양한 하이퍼파라미터가 있다.

그 중에서도 배치 사이즈(batch size)는 중요한 역할을 한다.

이번 포스팅에서는 배치 사이즈가 무엇이고 배치 사이즈를 작게 하거나 크게 했을 때의 효과를 예시 코드를 통해 알아보자.

그리고 배치 사이즈를 어떻게 선택해야하는지도 알아보자.

---

Batch size (배치 사이즈)

배치 사이즈는 머신러닝에서 모델을 훈련시킬 때 사용되는 데이터의 묶음 크기를 나타내는 개념이다.

훈련 데이터셋을 작은 묶음으로 나누어 모델에 순차적으로 제공하는 방식을 사용하는데, 이때 각 묶음의 크기가 바로 배치 사이즈이다.

 

배치 사이즈는 주로 확률적 경사 하강법(SGD) 및 미니 배치 경사 하강법과 같은 최적화 알고리즘에서 사용된다.

이 방법은 전체 데이터셋을 한 번에 처리하는 대신, 작은 배치 단위로 나누어 모델을 조금씩 업데이트하는 방식으로 효율적으로 학습한다.

 

예를 들어, 전체 데이터셋이 1000개의 샘플로 이루어져 있고 배치 사이즈가 64라면, 훈련 데이터를 64개의 샘플로 구성된 작은 배치로 나누어 모델에 제공한다.

따라서 모델은 이 작은 배치를 사용하여 가중치를 업데이트하고, 이 작업을 전체 데이터셋에 대해 반복하여 모델을 훈련시킨다.

 

배치 사이즈의 선택은 훈련 속도, 메모리 사용량, 모델의 수렴 안정성 등 다양한 요소를 고려하여 조절되어야 하며, 특정 문제와 데이터에 따라 최적의 값을 찾는 것이 중요하다.

---

배치 사이즈로 2의 거듭제곱을 주로 사용하는 이유

배치 사이즈를 2의 거듭제곱으로 주로 사용하는 이유에는 여러 가지 이론적인 이유와 실제적인 현상이 있다.

1. 하드웨어 최적화:
   많은 하드웨어 및 라이브러리에서는 2의 거듭제곱에 대한 최적화를 제공한다.
   특히 GPU와 같은 가속기는 2의 거듭제곱 배치 사이즈에 대해 더 효율적으로 동작하는 경우가 많다.
   이는 하드웨어의 병렬 처리 특성과 관련이 있다.
2. 메모리 최적화:
   많은 메모리 시스템이 크기가 2의 거듭제곱인 데이터에 대해 효율적으로 동작한다.
   이는 메모리의 구조와 관련이 있어서 2의 거듭제곱 배치 사이즈를 사용하면 데이터 메모리에 대한 최적화가 가능하다.
3. 코드 최적화 및 메모리 정렬:
   일부 최적화 기법 및 메모리 정렬 방식에서는 2의 거듭제곱이 더 효과적일 수 있다.
   특히 SIMD(Single Instruction, Multiple Data)와 같은 벡터 연산에서는 2의 거듭제곱이 더 효율적인 경우가 있다.
4. 하이퍼파라미터 튜닝의 편의성:
   하이퍼파라미터 튜닝 시, 배치 사이즈를 2의 거듭제곱으로 선택하면 다양한 배치 사이즈를 쉽게 시도해볼 수 있다. 또한, 일부 라이브러리는 2의 거듭제곱 배치 사이즈를 기본값으로 설정해 놓기도 한다.
5. 일관된 훈련 성능 관찰:
   2의 거듭제곱 배치 사이즈를 사용하면 다양한 크기의 배치 사이즈 간에 훈련 성능을 일관되게 비교할 수 있다.
   실험 결과를 논문이나 보고서로 전달할 때 통일성을 유지하는 데 도움이 된다.

결론적으로 컴퓨터가 2의 거듭제곱 개수의 데이터를 받아들이기 편하다는 이유가 주된 이유이다.

---

배치 사이즈 크기에 따른 효과

배치 사이즈에 따른 장단점을 알아보자.

 

배치 사이즈 키우기:

• 장점:
  • 훈련 속도 향상:
    큰 배치 사이즈를 사용하면 한 번에 더 많은 데이터를 처리할 수 있어 훈련 속도가 향상된다.
    특히 높은 병렬성을 지닌 하드웨어(GPU 등)에서 빠른 모델 업데이트가 가능하다.
  • 효율적인 하드웨어 활용:
    GPU와 같은 가속기는 대부분 큰 배치 사이즈에 대해 최적화되어 있어, 이를 활용하여 계산 자원을 효율적으로 사용할 수 있다.
• 단점:
  • 메모리 부담:
    큰 배치 사이즈는 더 많은 메모리를 필요로 하며, 이는 메모리 부담을 증가시킬 수 있다.
    그래서 메모리가 제한적인 환경에서는 주의가 필요하다.
  • 수렴 불안정성:
    특정 모델에서는 큰 배치 사이즈가 수렴에 어려움을 줄 수 있으며, 최적의 전역 최소값에 도달하기 어려울 수 있다.

 

배치 사이즈 줄이기:

• 장점:
  • 메모리 효율성:
    작은 배치 사이즈는 적은 메모리를 사용하며, 제한된 자원으로도 모델을 효과적으로 학습시킬 수 있다.
    특히 메모리가 한정된 환경에서 유용하다.
  • 미세한 패턴 학습:
    작은 배치 사이즈는 미세한 업데이트를 가능하게 하여 모델이 데이터의 미세한 패턴을 더 잘 학습할 수 있다.
• 단점:
  • 훈련 속도 감소:
    작은 배치 사이즈는 훈련 속도를 떨어뜨릴 수 있다.
    특히 GPU와 같은 가속기의 효율을 최대로 활용하기 어려워진다.
  • 노이즈와 수렴 불안정성:
    미세 패턴을 잘 학습한다는 이야기는 훈련 데이터의 노이즈에도 민감해진다는 뜻과도 일맥상통하다.
    따라서 모델이 학습하는 과정에서 수렴이 불안정해질 수 있다.

배치 사이즈를 키우거나 줄일 때, 각각의 장단점을 고려하여 특정 모델과 데이터에 가장 적합한 크기를 선택해야 한다.

---

적절한 배치 사이즈의 선택

배치 사이즈의 선택을 일단 결론부터 말하자면 많은 경험이 필요하다.

그래도 어떤 측면들을 고려할 수 있는지 살펴보자.

• 데이터셋의 크기와 품질:
  데이터셋이 크고 다양하다면 상대적으로 큰 배치 사이즈가 효과적일 수 있다.
  그러나 데이터가 제한적이거나 불균형하다면 작은 배치 사이즈가 미세한 패턴을 잘 학습할 수 있다.
• 계산 자원 및 하드웨어 제한:
  사용 가능한 계산 자원과 하드웨어 성능은 배치 사이즈 선택에 큰 영향을 미친다.
  GPU와 같은 디바이스를 효과적으로 활용하기 위해선 적절한 크기를 선택해야 한다.
• 모델의 복잡도:
  모델이 복잡하고 많은 파라미터를 가지면 상대적으로 큰 배치 사이즈를 사용하는 것이 수렴을 안정화시킬 수 있다.
  그러나 간단한 모델에서는 작은 배치 사이즈도 충분할 수 있다.
• 메모리 관리:
  사용 가능한 메모리 양도 고려해야 한다.
  메모리 부족으로 인한 예외 상황을 피하기 위해 적절한 크기를 선택해야 한다.
• 훈련 속도와 수렴 속도:
  큰 배치 사이즈는 훈련 속도를 높일 수 있지만, 수렴이 어려워질 수 있다.
  작은 배치 사이즈는 훈련 속도가 감소할 수 있지만, 모델이 세밀한 패턴을 더 잘 학습할 수 있다.
• 하이퍼파라미터 튜닝:
  배치 사이즈는 다른 하이퍼파라미터와 함께 튜닝되어야 한다.
  모델의 최적화를 위해 실험을 통해 최적의 값을 찾는 것이 중요하다.

따라서 배치 사이즈를 선택할 때에는 이러한 다양한 측면을 종합적으로 고려하여 모델과 데이터에 최적화된 크기를 찾아야 한다.

실험과 검증을 통해 최적의 배치 사이즈를 조정하는 것이 바람직하다.

---

PyTorch 예시 코드를 통해 살펴보는 배치 사이즈에 따른 학습 속도와 메모리 사용량 비교

MNIST 데이터셋으로 배치 사이즈에 따른 성능 차이를 비교해보는 코드를 만들어보고 직접 돌려보자.

 

먼저 필요한 라이브러리를 임포트한다.

import gc
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

이제 MNIST 데이터를 가져오자.

# 데이터 전처리 및 로딩
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

학습시킬 간단한 모델도 만들자.

간단하게 128개의 노드가 있는 한 개의 히든 레이어를 가진 뉴럴넷을 만들자.

# 모델 정의
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

이제 학습하는 코드와 검증하는 코드를 만들자.

Nvidia GPU를 사용하는 유저의 경우 코드를 그대로 쓰고 Apple Silicon을 사용하는 유저는 16번째 줄을 주석 처리하고 17번째 줄을 주석 해제하여 사용하자.

# 함수: 트레이닝 및 테스트
def train(model, train_loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    start_time = time.time()
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    end_time = time.time()
    elapsed_time = end_time - start_time
    memory_usage = torch.cuda.memory_allocated() / 1024 / 1024  # GPU 메모리 사용량 (MB)
    # memory_usage = torch.mps.current_allocated_memory() / 1024 / 1024  # 애플 실리콘 유저의 경우 윗 줄 대신 이 코드 사용
    return running_loss / len(train_loader), elapsed_time, memory_usage

def test(model, test_loader, criterion, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    accuracy = correct / total
    return accuracy

위 코드에서 torch.cuda.memory_allocated() 함수를 사용하여 GPU 메모리 사용량을 확인할 수 있다.

이 함수는 바이트 단위로 GPU에 할당된 메모리 크기를 반환한다.

다만, 메모리 사용량은 정확한 측정이 어렵기 때문에 시스템 모니터링 도구 등을 이용하는 것이 어쩌면 더 정확할 수도 있다.

 

이제 학습에 사용되는 하이퍼파라미터들을 정의하고 모델을 학습시켜보자.

배치 사이즈에 따른 학습에 걸리는 시간과 메모리의 사용량도 측정해보자.

# 트레이닝 파라미터 설정
batch_sizes = [32, 1024, 8192]
learning_rate = 0.001
epochs = 5
device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')

# 각 배치 사이즈에 대한 트레이닝 및 테스트
for batch_size in batch_sizes:

    # 배치 사이즈 변경 전에 메모리 정리
    torch.cuda.empty_cache()
    # torch.mps.empty_cache()  # 애플 실리콘 전용
    gc.collect()
    
    print(f"\nTraining with batch size: {batch_size}")
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    model = SimpleNN().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(epochs):
        train_loss, train_time, train_memory = train(model, train_loader, criterion, optimizer, device)
        accuracy = test(model, test_loader, criterion, device)
        print(f"Epoch {epoch + 1}/{epochs} => "
              f"Train Loss: {train_loss:.4f}, "
              f"Train Time: {train_time:.2f} seconds, "
              f"Train Memory: {train_memory:.2f} MB, "
              f"Test Accuracy: {accuracy:.4f}")

드라마틱한 변화를 보기 위해 배치 사이즈로는 다소 극단적으로 보이는 값인 32, 1024, 8192을 선택했다.

마찬가지로 Nvidia GPU를 사용하는 유저의 경우 코드를 그대로 쓰고 Apple Silicon을 사용하는 유저는 11번째 줄을 주석 처리하고 12번째 줄을 주석 해제하여 사용하자.

그리고 gc.collect()는 파이썬의 garbage collector를 수동으로 실행하여 불필요한 객체를 해제하는 함수로 새로운 배치 사이즈에서 정확한 메모리 측정을 위해 이전에 사용된 메모리를 정리하는 기능이다.

 

이제 각자 코드를 돌려보자.

내 결과는 아래와 같다.

Batch size에 따른 학습 시간과 메모리 사용량의 변화

위에서 설명한대로 배치 사이즈가 커질 수록 학습 시간이 줄어들고 메모리 사용량은 늘어나는 것을 확인할 수 있다.

---

살펴본 것처럼 배치 사이즈의 선택은 모델의 학습과 성능을 좌우한다.

결국 최적의 배치 사이즈 선택은 많은 경험에서 비롯되는 것이니 이것 저것 시도 해보자.