PyTorch를 사용하여 머신러닝 모델을 개발할 때는 데이터셋을 훈련, 검증, 테스트 세트로 나누는 것이 중요합니다. 이를 통해 모델의 성능을 평가하고, 과적합을 방지할 수 있습니다.

이번 포스트에서는 PyTorch에서 데이터셋을 나누는 방법에 대해 소개하겠습니다.

1. 데이터셋 불러오기

PyTorch에서 데이터셋을 사용하기 위해 먼저 데이터를 불러와야 합니다. 이를 위해 torchvision.datasets 모듈을 사용할 수 있습니다. 예를 들면, CIFAR-10 데이터셋을 불러오는 코드는 다음과 같습니다.

import torchvision.datasets as datasets

# CIFAR-10 데이터셋 불러오기
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True)

2. 데이터 세트를 나누기

PyTorch에서는 torch.utils.data.random_split 함수를 사용하여 데이터셋을 훈련, 검증, 테스트 세트로 나눌 수 있습니다. 이 함수는 입력으로 전체 데이터셋과 각 세트의 크기(비율)를 받습니다.

from torch.utils.data import random_split

# 데이터셋을 훈련, 검증, 테스트 세트로 나누기
train_size = int(0.8 * len(dataset))
val_size = int(0.1 * len(dataset))
test_size = len(dataset) - train_size - val_size

train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])

위의 예제에서는 전체 데이터셋을 80%의 훈련 세트, 10%의 검증 세트, 10%의 테스트 세트로 나누고 있습니다. random_split 함수는 무작위로 데이터를 섞고 분할합니다.

3. 데이터로더 생성하기

마지막으로 각 세트에 대한 데이터로더를 생성해야 합니다. 데이터로더는 모델에 데이터를 제공하기 위해 사용됩니다.

import torch.utils.data as data

# 데이터로더 생성하기
batch_size = 64

train_loader = data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = data.DataLoader(val_dataset, batch_size=batch_size)
test_loader = data.DataLoader(test_dataset, batch_size=batch_size)

위의 예제에서는 DataLoader 클래스를 사용하여 각 세트의 데이터로더를 생성하고 있습니다. batch_size는 한 번에 읽어들일 데이터의 개수입니다. shuffle=True로 설정하면 데이터가 학습 과정에서 섞여서 사용됩니다.

이제 데이터셋을 훈련, 검증, 테스트 세트로 나누는 방법에 대해 알아보았습니다. 이를 통해 PyTorch를 사용하여 데이터셋을 효율적으로 관리하고 모델을 훈련할 수 있습니다.

참고 자료

• PyTorch 공식 문서: torchvision.datasets
• PyTorch 공식 문서: torch.utils.data