PyTorch 모델 훈련 속도 개선 방법

PyTorch 모델 훈련 속도 개선 방법

딥러닝 모델을 PyTorch에서 학습할 때 훈련 속도를 최적화하는 것은 매우 중요한 과제입니다. 훈련 속도를 높이면 더 많은 실험을 수행할 수 있으며, 개발 과정에서 모델을 더 빠르게 개선할 수 있습니다. 본 글에서는 PyTorch에서 모델 훈련 속도를 향상시키기 위한 다양한 기법들을 다룹니다.

1. 데이터 로딩 최적화

훈련 속도를 높이기 위해서는 데이터 로딩 과정부터 최적화해야 합니다. PyTorch의 DataLoader를 활용하여 병렬 처리를 수행하면 데이터 로딩 속도를 향상시킬 수 있습니다.

1.1 num_workers 조정

DataLoader에서 num_workers 값을 증가시키면 데이터 로딩이 여러 프로세스에서 동시에 수행되므로 속도를 향상시킬 수 있습니다.

from torch.utils.data import DataLoader
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)

• num_workers: CPU 코어 수에 맞게 조정하여 성능을 극대화합니다.
• pin_memory=True: GPU로 데이터를 로드할 때 속도를 향상시킵니다.

2. GPU 활용 최적화

2.1 torch.cuda.amp를 활용한 Mixed Precision Training

PyTorch의 torch.cuda.amp를 활용하면 부동소수점 연산을 혼합 정밀도로 수행하여 속도를 높이고 메모리 사용량을 줄일 수 있습니다.

import torch

scaler = torch.cuda.amp.GradScaler()

def train_step(model, dataloader, loss_fn, optimizer, device):
    model.train()
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()

        with torch.cuda.amp.autocast():
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

2.2 torch.backends.cudnn.benchmark 설정

CUDA 백엔드에서 torch.backends.cudnn.benchmark = True를 설정하면 최적화된 커널을 자동으로 선택하여 성능을 향상시킬 수 있습니다.

import torch

torch.backends.cudnn.benchmark = True

3. 배치 크기 조정

배치 크기를 최대로 설정하면 GPU의 연산 효율을 높일 수 있습니다. 하지만 너무 크게 설정하면 OOM(Out of Memory) 오류가 발생할 수 있으므로 적절한 값을 찾아야 합니다.

3.1 배치 크기 자동 조정

아래 코드를 활용하면 메모리를 초과하지 않는 최대 배치 크기를 찾을 수 있습니다.

def find_max_batch_size(model, dataloader, loss_fn, optimizer, device):
    batch_size = 32
    while True:
        try:
            sample_data = next(iter(dataloader))
            inputs, targets = sample_data[0].to(device), sample_data[1].to(device)
            model(inputs)
            batch_size *= 2
        except RuntimeError:
            break
    return batch_size // 2

4. 연산 그래프 최적화

4.1 with torch.no_grad() 활용

모델 평가 시에는 그래디언트 계산이 필요 없으므로 torch.no_grad()를 사용하면 메모리 사용량을 줄이고 속도를 높일 수 있습니다.

with torch.no_grad():
    for inputs, targets in test_dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)

4.2 detach() 활용

불필요한 연산 그래프 생성을 방지하기 위해 detach()를 사용하면 성능을 더욱 최적화할 수 있습니다.

outputs = model(inputs).detach()

5. 효율적인 옵티마이저 사용

5.1 AdamW 옵티마이저 활용

Adam보다 AdamW 옵티마이저를 사용하면 학습 속도가 개선될 수 있습니다.

import torch.optim as optim

optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-2)

5.2 Learning Rate Scheduler 활용

학습률 스케줄링을 통해 훈련 속도를 개선할 수 있습니다.

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
for epoch in range(epochs):
    train_step(model, dataloader, loss_fn, optimizer, device)
    scheduler.step()

6. TorchScript를 활용한 모델 최적화

PyTorch의 TorchScript를 활용하면 모델을 컴파일하여 속도를 높일 수 있습니다.

traced_model = torch.jit.trace(model, torch.randn(1, 3, 224, 224).to(device))

7. 결론

위에서 소개한 기법들을 조합하여 활용하면 PyTorch에서 모델 훈련 속도를 효과적으로 개선할 수 있습니다. 데이터 로딩, GPU 활용, 배치 크기 조정, 연산 그래프 최적화, 효율적인 옵티마이저 사용 등을 고려하여 최적의 성능을 도출하는 것이 중요합니다. 실험을 통해 가장 적합한 설정을 찾아 적용해 보시기 바랍니다.